mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2026-10-30 17:00:48 +01:00
* vector: DeleteIndex format guard, sticky discovery, advertise normalize - handleDeleteIndex verifies the table carries the Lance format before DeleteTable — an Iceberg table under the same namespace is not an index and must not be removed by it. - discovery keeps the current worker while it is still registered — membership-order churn can no longer cancel in-flight calls, and the resolver closes each master conn inside a 5s deadline. - advertised_addr strips scheme and .grpcPort so the UDP probe reaches every supported master form. * vector: empty-follower retry, format check without read grant - discoverVectorWorkers keeps looking past an empty master reply — followers list nothing when workers registered with the leader. - the DeleteIndex format check uses the internal catalog read, so DeleteTable-only callers don't gain a GetTable requirement, and the format must be Lance exactly — format-less tables no longer pass.
1289 lines
43 KiB
Go
1289 lines
43 KiB
Go
package s3vectors
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/seaweedfs/seaweedfs/weed/glog"
|
|
"github.com/seaweedfs/seaweedfs/weed/pb/vector_pb"
|
|
"github.com/seaweedfs/seaweedfs/weed/s3api/s3tables"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
"google.golang.org/protobuf/types/known/structpb"
|
|
)
|
|
|
|
var handlers = map[string]func(*Server, http.ResponseWriter, *http.Request){
|
|
"CreateVectorBucket": (*Server).handleCreateVectorBucket,
|
|
"GetVectorBucket": (*Server).handleGetVectorBucket,
|
|
"ListVectorBuckets": (*Server).handleListVectorBuckets,
|
|
"DeleteVectorBucket": (*Server).handleDeleteVectorBucket,
|
|
"CreateIndex": (*Server).handleCreateIndex,
|
|
"GetIndex": (*Server).handleGetIndex,
|
|
"ListIndexes": (*Server).handleListIndexes,
|
|
"DeleteIndex": (*Server).handleDeleteIndex,
|
|
"PutVectors": (*Server).handlePutVectors,
|
|
"GetVectors": (*Server).handleGetVectors,
|
|
"ListVectors": (*Server).handleListVectors,
|
|
"DeleteVectors": (*Server).handleDeleteVectors,
|
|
"QueryVectors": (*Server).handleQueryVectors,
|
|
}
|
|
|
|
// lance table ids put indexes under a "vector" namespace inside their bucket.
|
|
const indexNamespace = "vector"
|
|
|
|
// Vector bucket records are s3tables table buckets; the S3 Vectors ARN form is
|
|
// arn:aws:s3vectors:{region}:{account}:bucket/{name}[/index/{name}].
|
|
func (s *Server) vectorBucketArn(bucket string) string {
|
|
return fmt.Sprintf("arn:aws:s3vectors:%s:%s:bucket/%s", s.region, s.accountID, bucket)
|
|
}
|
|
|
|
func (s *Server) indexArn(bucket, index string) string {
|
|
return fmt.Sprintf("%s/index/%s", s.vectorBucketArn(bucket), index)
|
|
}
|
|
|
|
// s3tablesArn renders the same bucket in the table-bucket ARN form the S3
|
|
// Tables catalog keys records by.
|
|
func (s *Server) s3tablesArn(bucket string) string {
|
|
return fmt.Sprintf("arn:aws:s3tables:%s:%s:bucket/%s", s.region, s.accountID, bucket)
|
|
}
|
|
|
|
func arnParts(v *vectorBucketARNs, name, arn string) error {
|
|
if arn == "" {
|
|
v.bucket = name
|
|
return nil
|
|
}
|
|
if !strings.HasPrefix(arn, "arn:") {
|
|
v.bucket = arn
|
|
return nil
|
|
}
|
|
i := strings.Index(arn, "bucket/")
|
|
if i < 0 {
|
|
return fmt.Errorf("malformed ARN %q", arn)
|
|
}
|
|
rest := arn[i+len("bucket/"):]
|
|
parts := strings.Split(rest, "/index/")
|
|
v.bucket = parts[0]
|
|
if v.bucket == "" {
|
|
return fmt.Errorf("malformed ARN %q", arn)
|
|
}
|
|
if len(parts) == 2 {
|
|
v.index = parts[1]
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// resolveIndexARNs fills bucket and index from any combination of plain names,
|
|
// a bucket ARN, and a full index ARN — callers may supply only an indexArn.
|
|
func resolveIndexARNs(bucketName, bucketArn, indexName, indexArn string) (vectorBucketARNs, error) {
|
|
var v vectorBucketARNs
|
|
if strings.HasPrefix(indexArn, "arn:") {
|
|
if err := arnParts(&v, "", indexArn); err != nil {
|
|
return v, err
|
|
}
|
|
} else if indexArn != "" {
|
|
return v, fmt.Errorf("indexArn must be an ARN, got %q", indexArn)
|
|
}
|
|
if bucketArn != "" {
|
|
if !strings.HasPrefix(bucketArn, "arn:") {
|
|
return v, fmt.Errorf("vectorBucketArn must be an ARN, got %q", bucketArn)
|
|
}
|
|
if err := arnParts(&v, "", bucketArn); err != nil {
|
|
return v, err
|
|
}
|
|
}
|
|
if v.bucket == "" {
|
|
v.bucket = bucketName
|
|
}
|
|
if v.index == "" {
|
|
v.index = indexName
|
|
}
|
|
if v.bucket == "" || v.index == "" {
|
|
return v, fmt.Errorf("vectorBucketName and indexName (or indexArn) are required")
|
|
}
|
|
return v, nil
|
|
}
|
|
|
|
// indexRecord fetches the index's catalog record; the call also authorizes
|
|
// the request identity against the index resource.
|
|
func (s *Server) indexRecord(ctx context.Context, r *http.Request, bucket, index string) (map[string]interface{}, error) {
|
|
var resp map[string]interface{}
|
|
err := s.catalog(ctx, "GetTable", &s3tables.GetTableRequest{
|
|
TableBucketARN: s.s3tablesArn(bucket),
|
|
Namespace: []string{indexNamespace},
|
|
Name: index,
|
|
}, &resp, r)
|
|
return resp, err
|
|
}
|
|
|
|
func (s *Server) catalog(ctx context.Context, operation string, req, resp interface{}, r *http.Request) error {
|
|
return s.tables.Execute(ctx, s.filerClient, operation, req, resp, identityOf(r))
|
|
}
|
|
|
|
// internalCatalog runs catalog bookkeeping under the trusted identity. It is
|
|
// for reads and writes that accompany — never replace — the operation's own
|
|
// permission check: bucket format lookups, the implicit index namespace, and
|
|
// rollback of a record this request just created.
|
|
func (s *Server) internalCatalog(ctx context.Context, operation string, req, resp interface{}) error {
|
|
return s.internal.Execute(ctx, s.filerClient, operation, req, resp, "")
|
|
}
|
|
|
|
func catalogErrCode(err error) string {
|
|
var apiErr *s3tables.S3TablesError
|
|
if errors.As(err, &apiErr) {
|
|
return apiErr.Type
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func catalogErr(w http.ResponseWriter, err error) {
|
|
var apiErr *s3tables.S3TablesError
|
|
if !errors.As(err, &apiErr) {
|
|
writeError(w, http.StatusInternalServerError, "InternalFailure", "internal error")
|
|
return
|
|
}
|
|
switch {
|
|
case apiErr.Type == s3tables.ErrCodeAccessDenied:
|
|
writeError(w, http.StatusForbidden, "AccessDeniedException", "access denied")
|
|
case strings.HasPrefix(apiErr.Type, "NoSuch"):
|
|
writeError(w, http.StatusNotFound, "NotFoundException", apiErr.Message)
|
|
case strings.Contains(apiErr.Type, "AlreadyExists") || strings.Contains(apiErr.Type, "NotEmpty") ||
|
|
apiErr.Type == s3tables.ErrCodeConflict:
|
|
writeError(w, http.StatusConflict, "ConflictException", apiErr.Message)
|
|
case apiErr.Type == s3tables.ErrCodeInternalError:
|
|
writeError(w, http.StatusInternalServerError, "InternalFailure", apiErr.Message)
|
|
default:
|
|
writeError(w, http.StatusBadRequest, "ValidationException", apiErr.Message)
|
|
}
|
|
}
|
|
|
|
// vectorErr maps a VectorService gRPC status onto the API's error shape.
|
|
// Engine messages reach clients only for codes that describe the request
|
|
// itself; transport and internal details stay in the server log.
|
|
func vectorErr(w http.ResponseWriter, err error) {
|
|
msg := status.Convert(err).Message()
|
|
switch status.Code(err) {
|
|
case codes.NotFound:
|
|
writeError(w, http.StatusNotFound, "NotFoundException", msg)
|
|
case codes.InvalidArgument, codes.FailedPrecondition, codes.OutOfRange:
|
|
writeError(w, http.StatusBadRequest, "ValidationException", msg)
|
|
case codes.PermissionDenied, codes.Unauthenticated:
|
|
writeError(w, http.StatusForbidden, "AccessDeniedException", "access denied")
|
|
case codes.Unavailable:
|
|
writeError(w, http.StatusServiceUnavailable, "ServiceUnavailableException", "vector service unavailable")
|
|
default:
|
|
glog.V(1).Infof("S3 Vectors: engine error: %v", err)
|
|
writeError(w, http.StatusInternalServerError, "InternalFailure", "internal error")
|
|
}
|
|
}
|
|
|
|
// lanceBucket reports whether bucket is a table bucket that can hold vector
|
|
// indexes: a declared Lance bucket, or one made before formats existed. An
|
|
// explicit ICEBERG bucket is not a vector bucket and reads as absent. The
|
|
// lookup is internal so the call needs no read grant — the requested
|
|
// operation's own permission still gates it.
|
|
func (s *Server) lanceBucket(ctx context.Context, bucket string) error {
|
|
var resp struct {
|
|
Format string `json:"format"`
|
|
}
|
|
err := s.internalCatalog(ctx, "GetTableBucket",
|
|
&s3tables.GetTableBucketRequest{TableBucketARN: s.s3tablesArn(bucket)}, &resp)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if resp.Format != "" && resp.Format != s3tables.FormatLance {
|
|
return &s3tables.S3TablesError{
|
|
Type: s3tables.ErrCodeNoSuchBucket,
|
|
Message: fmt.Sprintf("vector bucket %s not found", bucket),
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ensureIndexNamespace declares the implicit index namespace when absent. The
|
|
// calls are internal: a caller's CreateTable grant is what authorizes the
|
|
// index, so it must not also need namespace reads or writes.
|
|
func (s *Server) ensureIndexNamespace(ctx context.Context, bucket string) error {
|
|
var resp map[string]interface{}
|
|
err := s.internalCatalog(ctx, "GetNamespace", &s3tables.GetNamespaceRequest{
|
|
TableBucketARN: s.s3tablesArn(bucket),
|
|
Namespace: []string{indexNamespace},
|
|
}, &resp)
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
if code := catalogErrCode(err); code != s3tables.ErrCodeNoSuchNamespace && code != s3tables.ErrCodeNoSuchBucket {
|
|
return err
|
|
}
|
|
err = s.internalCatalog(ctx, "CreateNamespace", &s3tables.CreateNamespaceRequest{
|
|
TableBucketARN: s.s3tablesArn(bucket),
|
|
Namespace: []string{indexNamespace},
|
|
}, &resp)
|
|
// A concurrent CreateIndex may have won the create.
|
|
if err != nil && catalogErrCode(err) != s3tables.ErrCodeNamespaceAlreadyExists {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// authorizeIndexWrite requires the update grant on the index's catalog table
|
|
// — the same check credential vending uses — so a read-only GetTable grant
|
|
// never reaches a mutation.
|
|
func (s *Server) authorizeIndexWrite(ctx context.Context, r *http.Request, bucket, index string) error {
|
|
return s.tables.AuthorizeUpdateTable(ctx, s.filerClient, &s3tables.UpdateTableRequest{
|
|
TableBucketARN: s.s3tablesArn(bucket),
|
|
Namespace: []string{indexNamespace},
|
|
Name: index,
|
|
}, identityOf(r))
|
|
}
|
|
|
|
func (s *Server) handleCreateVectorBucket(w http.ResponseWriter, r *http.Request) {
|
|
var req CreateVectorBucketRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
if req.VectorBucketName == "" {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "vectorBucketName is required")
|
|
return
|
|
}
|
|
// The catalog has no encryption boundary to honor the configuration
|
|
// against; accepting it would silently promise semantics we do not keep.
|
|
if len(req.EncryptionConfiguration) > 0 {
|
|
writeError(w, http.StatusBadRequest, "ValidationException",
|
|
"encryptionConfiguration is not supported")
|
|
return
|
|
}
|
|
var resp map[string]interface{}
|
|
err := s.catalog(r.Context(), "CreateTableBucket",
|
|
&s3tables.CreateTableBucketRequest{Name: req.VectorBucketName, Format: s3tables.FormatLance}, &resp, r)
|
|
if err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, &CreateVectorBucketResponse{VectorBucketArn: s.vectorBucketArn(req.VectorBucketName)})
|
|
}
|
|
|
|
func (s *Server) handleGetVectorBucket(w http.ResponseWriter, r *http.Request) {
|
|
var req GetVectorBucketRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
var arns vectorBucketARNs
|
|
if err := arnParts(&arns, req.VectorBucketName, req.VectorBucketArn); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
var resp map[string]interface{}
|
|
err := s.catalog(r.Context(), "GetTableBucket",
|
|
&s3tables.GetTableBucketRequest{TableBucketARN: s.s3tablesArn(arns.bucket)}, &resp, r)
|
|
if err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
// An explicit ICEBERG bucket is not a vector bucket.
|
|
if format, _ := resp["format"].(string); format != "" && format != s3tables.FormatLance {
|
|
writeError(w, http.StatusNotFound, "NotFoundException",
|
|
fmt.Sprintf("vector bucket %s not found", arns.bucket))
|
|
return
|
|
}
|
|
created, _ := time.Parse(time.RFC3339, fmt.Sprint(resp["createdAt"]))
|
|
writeJSON(w, http.StatusOK, &GetVectorBucketResponse{VectorBucket: &VectorBucket{
|
|
VectorBucketArn: s.vectorBucketArn(arns.bucket),
|
|
VectorBucketName: arns.bucket,
|
|
CreationTime: created,
|
|
}})
|
|
}
|
|
|
|
func (s *Server) handleListVectorBuckets(w http.ResponseWriter, r *http.Request) {
|
|
var req ListVectorBucketsRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
var resp struct {
|
|
TableBuckets []struct {
|
|
Arn string `json:"arn"`
|
|
Name string `json:"name"`
|
|
CreatedAt string `json:"createdAt"`
|
|
Format string `json:"format"`
|
|
} `json:"tableBuckets"`
|
|
ContinuationToken string `json:"continuationToken"`
|
|
}
|
|
err := s.catalog(r.Context(), "ListTableBuckets",
|
|
&s3tables.ListTableBucketsRequest{
|
|
MaxBuckets: req.MaxResults, ContinuationToken: req.NextToken, Prefix: req.Prefix,
|
|
}, &resp, r)
|
|
if err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
// The catalog token continues the filtered page: ICEBERG buckets drop out
|
|
// of the listing, Lance and format-less legacy buckets stay.
|
|
out := &ListVectorBucketsResponse{NextToken: resp.ContinuationToken}
|
|
for _, b := range resp.TableBuckets {
|
|
if b.Format != "" && b.Format != s3tables.FormatLance {
|
|
continue
|
|
}
|
|
created, _ := time.Parse(time.RFC3339, b.CreatedAt)
|
|
out.VectorBuckets = append(out.VectorBuckets, &VectorBucket{
|
|
VectorBucketArn: s.vectorBucketArn(b.Name),
|
|
VectorBucketName: b.Name,
|
|
CreationTime: created,
|
|
})
|
|
}
|
|
writeJSON(w, http.StatusOK, out)
|
|
}
|
|
|
|
func (s *Server) handleDeleteVectorBucket(w http.ResponseWriter, r *http.Request) {
|
|
var req DeleteVectorBucketRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
var arns vectorBucketARNs
|
|
if err := arnParts(&arns, req.VectorBucketName, req.VectorBucketArn); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
// The format check needs no read grant — the delete below is the
|
|
// authorization boundary — but an ICEBERG bucket is not a vector bucket.
|
|
if err := s.lanceBucket(r.Context(), arns.bucket); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
var resp map[string]interface{}
|
|
err := s.catalog(r.Context(), "DeleteTableBucket",
|
|
&s3tables.DeleteTableBucketRequest{TableBucketARN: s.s3tablesArn(arns.bucket)}, &resp, r)
|
|
if err != nil && catalogErrCode(err) == s3tables.ErrCodeBucketNotEmpty {
|
|
// The implicit index namespace counts as bucket content even when it
|
|
// holds no indexes; drop it under the delete's own authorization and
|
|
// retry once. A namespace still holding indexes keeps the conflict.
|
|
var nsResp map[string]interface{}
|
|
nsErr := s.internalCatalog(context.WithoutCancel(r.Context()), "DeleteNamespace",
|
|
&s3tables.DeleteNamespaceRequest{
|
|
TableBucketARN: s.s3tablesArn(arns.bucket),
|
|
Namespace: []string{indexNamespace},
|
|
}, &nsResp)
|
|
if nsErr == nil {
|
|
err = s.catalog(r.Context(), "DeleteTableBucket",
|
|
&s3tables.DeleteTableBucketRequest{TableBucketARN: s.s3tablesArn(arns.bucket)}, &resp, r)
|
|
} else if code := catalogErrCode(nsErr); code != s3tables.ErrCodeNoSuchNamespace {
|
|
glog.V(1).Infof("S3 Vectors: drop internal namespace in %s: %v", arns.bucket, nsErr)
|
|
}
|
|
}
|
|
if err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, &DeleteVectorBucketResponse{})
|
|
}
|
|
|
|
func distanceMetric(metric string) (vector_pb.DistanceMetric, bool) {
|
|
switch strings.ToLower(metric) {
|
|
case "", "euclidean":
|
|
return vector_pb.DistanceMetric_L2, true
|
|
case "cosine":
|
|
return vector_pb.DistanceMetric_COSINE, true
|
|
case "dotproduct":
|
|
return vector_pb.DistanceMetric_DOT_PRODUCT, true
|
|
}
|
|
return vector_pb.DistanceMetric_L2, false
|
|
}
|
|
|
|
func tableID(bucket, index string) *vector_pb.Table {
|
|
return &vector_pb.Table{Id: []string{bucket, indexNamespace, index}}
|
|
}
|
|
|
|
func (s *Server) handleCreateIndex(w http.ResponseWriter, r *http.Request) {
|
|
var req CreateIndexRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
if strings.ToLower(req.DataType) != "float32" {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "only float32 vectors are supported")
|
|
return
|
|
}
|
|
metric, ok := distanceMetric(req.DistanceMetric)
|
|
if !ok {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "unknown distanceMetric "+req.DistanceMetric)
|
|
return
|
|
}
|
|
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
if req.Dimension <= 0 || req.Dimension > 4096 {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "dimension must be 1..=4096")
|
|
return
|
|
}
|
|
if err := s.lanceBucket(r.Context(), arns.bucket); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
// The create grant is exercised before any write: a denied caller must not
|
|
// leave even the implicit namespace behind. The namespace may not exist
|
|
// yet, so the check tolerates its absence.
|
|
if err := s.tables.AuthorizeCreateTable(r.Context(), s.filerClient, &s3tables.CreateTableRequest{
|
|
TableBucketARN: s.s3tablesArn(arns.bucket),
|
|
Namespace: []string{indexNamespace},
|
|
Name: arns.index,
|
|
AllowMissingNamespace: true,
|
|
}, identityOf(r)); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
// The namespace is implicit on AWS and internal bookkeeping here.
|
|
if err := s.ensureIndexNamespace(r.Context(), arns.bucket); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
// Only a caller who passed the create check may learn an index exists.
|
|
var existing map[string]interface{}
|
|
getErr := s.internalCatalog(r.Context(), "GetTable", &s3tables.GetTableRequest{
|
|
TableBucketARN: s.s3tablesArn(arns.bucket),
|
|
Namespace: []string{indexNamespace},
|
|
Name: arns.index,
|
|
}, &existing)
|
|
switch {
|
|
case getErr == nil:
|
|
writeError(w, http.StatusConflict, "ConflictException",
|
|
fmt.Sprintf("index %s already exists", arns.index))
|
|
return
|
|
case catalogErrCode(getErr) != s3tables.ErrCodeNoSuchTable:
|
|
catalogErr(w, getErr)
|
|
return
|
|
}
|
|
metadata := map[string]string{
|
|
"dimension": fmt.Sprint(req.Dimension),
|
|
"distanceMetric": req.DistanceMetric,
|
|
// The nonce marks the record this request created, so a rollback after
|
|
// an engine failure can never erase a pre-existing index.
|
|
"createNonce": uuid.NewString(),
|
|
}
|
|
var nonFilterable []string
|
|
if req.MetadataConfiguration != nil {
|
|
nonFilterable = req.MetadataConfiguration.NonFilterableMetadataKeys
|
|
if keys, err := json.Marshal(nonFilterable); err == nil {
|
|
metadata["nonFilterableMetadataKeys"] = string(keys)
|
|
}
|
|
}
|
|
metadataJSON, _ := json.Marshal(metadata)
|
|
err = s.catalog(r.Context(), "CreateTable", &s3tables.CreateTableRequest{
|
|
TableBucketARN: s.s3tablesArn(arns.bucket),
|
|
Namespace: []string{indexNamespace},
|
|
Name: arns.index,
|
|
Format: s3tables.FormatLance,
|
|
Metadata: &s3tables.TableMetadata{FullMetadata: metadataJSON},
|
|
}, nil, r)
|
|
if err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
_, err = s.vectors.CreateIndex(r.Context(), &vector_pb.CreateIndexRequest{
|
|
Table: tableID(arns.bucket, arns.index),
|
|
Dimension: uint32(req.Dimension),
|
|
DistanceMetric: metric,
|
|
NonFilterableMetadataKeys: nonFilterable,
|
|
})
|
|
if err != nil {
|
|
// Roll the index record back so a half-created index does not linger.
|
|
// The delete runs past a client disconnect and only removes the record
|
|
// this request created: the nonce must match, and the versionToken
|
|
// compare-and-swap keeps a racing create from being erased.
|
|
ctx := context.WithoutCancel(r.Context())
|
|
var rec map[string]interface{}
|
|
if gerr := s.internalCatalog(ctx, "GetTable", &s3tables.GetTableRequest{
|
|
TableBucketARN: s.s3tablesArn(arns.bucket),
|
|
Namespace: []string{indexNamespace},
|
|
Name: arns.index,
|
|
}, &rec); gerr == nil && recordNonce(rec) == metadata["createNonce"] {
|
|
var delResp map[string]interface{}
|
|
if derr := s.internalCatalog(ctx, "DeleteTable", &s3tables.DeleteTableRequest{
|
|
TableBucketARN: s.s3tablesArn(arns.bucket),
|
|
Namespace: []string{indexNamespace},
|
|
Name: arns.index,
|
|
VersionToken: recordVersionToken(rec),
|
|
}, &delResp); derr != nil {
|
|
glog.Errorf("S3 Vectors: rollback of catalog table %s/%s failed: %v", arns.bucket, arns.index, derr)
|
|
}
|
|
} else if gerr != nil {
|
|
glog.Errorf("S3 Vectors: rollback lookup for %s/%s failed: %v", arns.bucket, arns.index, gerr)
|
|
}
|
|
vectorErr(w, err)
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, &CreateIndexResponse{IndexArn: s.indexArn(arns.bucket, arns.index)})
|
|
}
|
|
|
|
// recordNonce and recordVersionToken read the catalog record's internal
|
|
// bookkeeping for the create-index rollback.
|
|
func recordNonce(record map[string]interface{}) string {
|
|
if meta, ok := record["metadata"].(map[string]interface{}); ok {
|
|
if fm, ok := meta["fullMetadata"].(map[string]interface{}); ok {
|
|
return fmt.Sprint(fm["createNonce"])
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func recordVersionToken(record map[string]interface{}) string {
|
|
token, _ := record["versionToken"].(string)
|
|
return token
|
|
}
|
|
|
|
func indexFromRecord(s *Server, bucket, name string, m map[string]interface{}) *Index {
|
|
idx := &Index{
|
|
IndexArn: s.indexArn(bucket, name),
|
|
IndexName: name,
|
|
VectorBucketName: bucket,
|
|
DataType: "float32",
|
|
}
|
|
if created, ok := m["createdAt"].(string); ok {
|
|
idx.CreationTime, _ = time.Parse(time.RFC3339, created)
|
|
}
|
|
if meta, ok := m["metadata"].(map[string]interface{}); ok {
|
|
if fm, ok := meta["fullMetadata"].(map[string]interface{}); ok {
|
|
var dim float64
|
|
fmt.Sscan(fmt.Sprint(fm["dimension"]), &dim)
|
|
idx.Dimension = int32(dim)
|
|
idx.DistanceMetric = fmt.Sprint(fm["distanceMetric"])
|
|
switch keys := fm["nonFilterableMetadataKeys"].(type) {
|
|
case string:
|
|
var list []string
|
|
if err := json.Unmarshal([]byte(keys), &list); err == nil && list != nil {
|
|
idx.MetadataConfiguration = &MetadataConfiguration{NonFilterableMetadataKeys: list}
|
|
}
|
|
case []interface{}:
|
|
var list []string
|
|
for _, k := range keys {
|
|
list = append(list, fmt.Sprint(k))
|
|
}
|
|
idx.MetadataConfiguration = &MetadataConfiguration{NonFilterableMetadataKeys: list}
|
|
}
|
|
}
|
|
}
|
|
if idx.DistanceMetric == "" || idx.DistanceMetric == "<nil>" {
|
|
idx.DistanceMetric = "euclidean"
|
|
}
|
|
return idx
|
|
}
|
|
|
|
func (s *Server) handleGetIndex(w http.ResponseWriter, r *http.Request) {
|
|
var req GetIndexRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
var resp map[string]interface{}
|
|
err = s.catalog(r.Context(), "GetTable", &s3tables.GetTableRequest{
|
|
TableBucketARN: s.s3tablesArn(arns.bucket),
|
|
Namespace: []string{indexNamespace},
|
|
Name: arns.index,
|
|
}, &resp, r)
|
|
if err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, &GetIndexResponse{Index: indexFromRecord(s, arns.bucket, arns.index, resp)})
|
|
}
|
|
|
|
func (s *Server) handleListIndexes(w http.ResponseWriter, r *http.Request) {
|
|
var req ListIndexesRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
var arns vectorBucketARNs
|
|
if err := arnParts(&arns, req.VectorBucketName, req.VectorBucketArn); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
var resp struct {
|
|
TableSummaries []struct {
|
|
Name string `json:"name"`
|
|
Format string `json:"format"`
|
|
CreatedAt string `json:"createdAt"`
|
|
} `json:"tables"`
|
|
ContinuationToken string `json:"continuationToken"`
|
|
}
|
|
err := s.catalog(r.Context(), "ListTables", &s3tables.ListTablesRequest{
|
|
TableBucketARN: s.s3tablesArn(arns.bucket),
|
|
Namespace: []string{indexNamespace},
|
|
MaxTables: req.MaxResults,
|
|
ContinuationToken: req.NextToken,
|
|
Prefix: req.Prefix,
|
|
}, &resp, r)
|
|
if err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
// Summaries carry only what the catalog listing knows; a summary that
|
|
// asserted a dimension or metric without a record lookup would lie.
|
|
out := &ListIndexesResponse{NextToken: resp.ContinuationToken}
|
|
for _, t := range resp.TableSummaries {
|
|
// A non-Lance table in the index namespace is not an index.
|
|
if t.Format != "" && t.Format != s3tables.FormatLance {
|
|
continue
|
|
}
|
|
created, _ := time.Parse(time.RFC3339, t.CreatedAt)
|
|
out.Indexes = append(out.Indexes, &IndexSummary{
|
|
IndexName: t.Name,
|
|
IndexArn: s.indexArn(arns.bucket, t.Name),
|
|
VectorBucketName: arns.bucket,
|
|
CreationTime: created,
|
|
})
|
|
}
|
|
writeJSON(w, http.StatusOK, out)
|
|
}
|
|
|
|
func (s *Server) handleDeleteIndex(w http.ResponseWriter, r *http.Request) {
|
|
var req DeleteIndexRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
// A named table is an index only while it carries the Lance format —
|
|
// an Iceberg or format-less table under the same namespace is not an
|
|
// index and must not be deleted by this operation.
|
|
// The format check needs no read grant — the delete below is the
|
|
// authorization boundary — and only a Lance-formatted record is a
|
|
// vector index.
|
|
var record map[string]interface{}
|
|
if err := s.internalCatalog(r.Context(), "GetTable", &s3tables.GetTableRequest{
|
|
TableBucketARN: s.s3tablesArn(arns.bucket),
|
|
Namespace: []string{indexNamespace},
|
|
Name: arns.index,
|
|
}, &record); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
if format, _ := record["format"].(string); format != s3tables.FormatLance {
|
|
writeError(w, http.StatusBadRequest, "ValidationException",
|
|
fmt.Sprintf("index %q is not a vector index", arns.index))
|
|
return
|
|
}
|
|
// Deleting the catalog record is the delete: the engine derives an index's
|
|
// existence from its table, and DropIndex is unimplemented on the shared
|
|
// service.
|
|
var resp map[string]interface{}
|
|
err = s.catalog(r.Context(), "DeleteTable", &s3tables.DeleteTableRequest{
|
|
TableBucketARN: s.s3tablesArn(arns.bucket),
|
|
Namespace: []string{indexNamespace},
|
|
Name: arns.index,
|
|
}, &resp, r)
|
|
if err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, &DeleteIndexResponse{})
|
|
}
|
|
|
|
func metadataStruct(m map[string]interface{}) (*structpb.Struct, error) {
|
|
if m == nil {
|
|
return nil, nil
|
|
}
|
|
return structpb.NewStruct(m)
|
|
}
|
|
|
|
func structToMap(s *structpb.Struct) map[string]interface{} {
|
|
if s == nil {
|
|
return nil
|
|
}
|
|
return s.AsMap()
|
|
}
|
|
|
|
func (s *Server) handlePutVectors(w http.ResponseWriter, r *http.Request) {
|
|
var req PutVectorsRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
if _, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
// A read grant on the index does not carry a write: mutations require the
|
|
// catalog's update grant, the same check credential vending applies.
|
|
if err := s.authorizeIndexWrite(r.Context(), r, arns.bucket, arns.index); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
vectors := make([]*vector_pb.Vector, 0, len(req.Vectors))
|
|
for _, v := range req.Vectors {
|
|
if v == nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "vectors contains a null element")
|
|
return
|
|
}
|
|
if v.Key == "" {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "vector key must not be empty")
|
|
return
|
|
}
|
|
metadata, err := metadataStruct(v.Metadata)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
var data []float32
|
|
if v.Data != nil {
|
|
data = v.Data.Float32
|
|
}
|
|
vectors = append(vectors, &vector_pb.Vector{Key: v.Key, Data: data, Metadata: metadata})
|
|
}
|
|
_, err = s.vectors.PutVectors(r.Context(), &vector_pb.PutVectorsRequest{
|
|
Table: tableID(arns.bucket, arns.index),
|
|
Vectors: vectors,
|
|
})
|
|
if err != nil {
|
|
vectorErr(w, err)
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, &PutVectorsResponse{})
|
|
}
|
|
|
|
func toOutputVectors(vectors []*vector_pb.Vector, returnData, returnMetadata bool) []*OutputVector {
|
|
out := make([]*OutputVector, 0, len(vectors))
|
|
for _, v := range vectors {
|
|
o := &OutputVector{Key: v.Key}
|
|
if returnData {
|
|
o.Data = &VectorData{Float32: v.Data}
|
|
}
|
|
if returnMetadata {
|
|
o.Metadata = structToMap(v.Metadata)
|
|
}
|
|
out = append(out, o)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (s *Server) handleGetVectors(w http.ResponseWriter, r *http.Request) {
|
|
var req GetVectorsRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
if _, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
resp, err := s.vectors.GetVectors(r.Context(), &vector_pb.GetVectorsRequest{
|
|
Table: tableID(arns.bucket, arns.index),
|
|
Keys: req.Keys,
|
|
ReturnData: req.ReturnData,
|
|
ReturnMetadata: req.ReturnMetadata,
|
|
})
|
|
if err != nil {
|
|
vectorErr(w, err)
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, &GetVectorsResponse{
|
|
Vectors: toOutputVectors(resp.Vectors, req.ReturnData, req.ReturnMetadata),
|
|
})
|
|
}
|
|
|
|
func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) {
|
|
var req ListVectorsRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
if _, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
if req.MaxResults < 0 {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "maxResults must be non-negative")
|
|
return
|
|
}
|
|
// AWS exposes a static split: segmentIndex/segmentCount divides the full
|
|
// listing. The engine service does not know AWS segments, so page through
|
|
// the index here and keep only the caller's share — O(index) per call,
|
|
// acceptable for the MVP. The engine's next_token is a row offset, so the
|
|
// caller's token is the offset to resume scanning from.
|
|
if req.SegmentCount > 0 {
|
|
if req.SegmentIndex < 0 || req.SegmentIndex >= req.SegmentCount {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "segmentIndex must be in [0, segmentCount)")
|
|
return
|
|
}
|
|
var offset uint64
|
|
if req.NextToken != "" {
|
|
var err error
|
|
if offset, err = strconv.ParseUint(req.NextToken, 10, 64); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "invalid nextToken")
|
|
return
|
|
}
|
|
}
|
|
out, nextToken, err := s.listSegmentVectors(r.Context(), tableID(arns.bucket, arns.index),
|
|
offset, req.MaxResults, uint64(req.SegmentCount), uint64(req.SegmentIndex), req.ReturnData, req.ReturnMetadata)
|
|
if err != nil {
|
|
vectorErr(w, err)
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, &ListVectorsResponse{
|
|
NextToken: nextToken,
|
|
Vectors: toOutputVectors(out, req.ReturnData, req.ReturnMetadata),
|
|
})
|
|
return
|
|
}
|
|
resp, err := s.vectors.ListVectors(r.Context(), &vector_pb.ListVectorsRequest{
|
|
Table: tableID(arns.bucket, arns.index),
|
|
MaxResults: uint32(req.MaxResults),
|
|
NextToken: req.NextToken,
|
|
ReturnData: req.ReturnData,
|
|
ReturnMetadata: req.ReturnMetadata,
|
|
})
|
|
if err != nil {
|
|
vectorErr(w, err)
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, &ListVectorsResponse{
|
|
NextToken: resp.NextToken,
|
|
Vectors: toOutputVectors(resp.Vectors, req.ReturnData, req.ReturnMetadata),
|
|
})
|
|
}
|
|
|
|
// Mirrors the engine's page sizing so a segmented listing pages like a plain
|
|
// one does.
|
|
const (
|
|
vectorsDefaultPage = 500
|
|
vectorsMaxPage = 1000
|
|
)
|
|
|
|
// listSegmentVectors scans the index from offset, keeping the rows that fall
|
|
// in the caller's segment until limit are collected or the index runs out.
|
|
// The engine's next_token is a row offset, so the returned token is the offset
|
|
// scanning resumes from.
|
|
func (s *Server) listSegmentVectors(ctx context.Context, table *vector_pb.Table, offset uint64, maxResults int, segmentCount, segmentIndex uint64, returnData, returnMetadata bool) ([]*vector_pb.Vector, string, error) {
|
|
limit := maxResults
|
|
if limit <= 0 {
|
|
limit = vectorsDefaultPage
|
|
} else if limit > vectorsMaxPage {
|
|
limit = vectorsMaxPage
|
|
}
|
|
var out []*vector_pb.Vector
|
|
nextToken := ""
|
|
token := strconv.FormatUint(offset, 10)
|
|
for {
|
|
page, err := s.vectors.ListVectors(ctx, &vector_pb.ListVectorsRequest{
|
|
Table: table,
|
|
MaxResults: vectorsMaxPage,
|
|
NextToken: token,
|
|
ReturnData: returnData,
|
|
ReturnMetadata: returnMetadata,
|
|
})
|
|
if err != nil {
|
|
return nil, "", err
|
|
}
|
|
base := offset
|
|
offset += uint64(len(page.Vectors))
|
|
full := false
|
|
for j, v := range page.Vectors {
|
|
if (base+uint64(j))%segmentCount == segmentIndex {
|
|
out = append(out, v)
|
|
if len(out) == limit {
|
|
if j+1 < len(page.Vectors) || page.NextToken != "" {
|
|
nextToken = strconv.FormatUint(base+uint64(j)+1, 10)
|
|
}
|
|
full = true
|
|
break
|
|
}
|
|
}
|
|
}
|
|
if full || len(page.Vectors) == 0 || page.NextToken == "" {
|
|
break
|
|
}
|
|
token = page.NextToken
|
|
}
|
|
return out, nextToken, nil
|
|
}
|
|
|
|
func (s *Server) handleDeleteVectors(w http.ResponseWriter, r *http.Request) {
|
|
var req DeleteVectorsRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
if _, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
if err := s.authorizeIndexWrite(r.Context(), r, arns.bucket, arns.index); err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
_, err = s.vectors.DeleteVectors(r.Context(), &vector_pb.DeleteVectorsRequest{
|
|
Table: tableID(arns.bucket, arns.index),
|
|
Keys: req.Keys,
|
|
})
|
|
if err != nil {
|
|
vectorErr(w, err)
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, &DeleteVectorsResponse{})
|
|
}
|
|
|
|
// translateFilter renders an AWS S3 Vectors filter document as a Lance SQL
|
|
// predicate over the metadata JSON column. Only the documented operators are
|
|
// translated; anything else fails loudly instead of silently overmatching.
|
|
func translateFilter(filter map[string]interface{}) (string, error) {
|
|
if len(filter) == 0 {
|
|
return "", nil
|
|
}
|
|
return filterExpr(filter)
|
|
}
|
|
|
|
// nonFilterableField returns the first blocked metadata key a filter
|
|
// references; the bool reports whether any matched.
|
|
func nonFilterableField(cfg *MetadataConfiguration, filter map[string]interface{}) (string, bool) {
|
|
if cfg == nil {
|
|
return "", false
|
|
}
|
|
fields := map[string]bool{}
|
|
filterFields(filter, fields)
|
|
for _, k := range cfg.NonFilterableMetadataKeys {
|
|
if fields[k] {
|
|
return k, true
|
|
}
|
|
}
|
|
return "", false
|
|
}
|
|
|
|
// filterFields collects the field names a filter expression references —
|
|
// $-operators recurse; every other key is a field.
|
|
func filterFields(filter map[string]interface{}, out map[string]bool) {
|
|
for key, cond := range filter {
|
|
if strings.HasPrefix(key, "$") {
|
|
if sub, ok := cond.([]interface{}); ok {
|
|
for _, item := range sub {
|
|
if m, ok := item.(map[string]interface{}); ok {
|
|
filterFields(m, out)
|
|
}
|
|
}
|
|
}
|
|
continue
|
|
}
|
|
out[key] = true
|
|
}
|
|
}
|
|
|
|
func filterExpr(filter map[string]interface{}) (string, error) {
|
|
parts := make([]string, 0, len(filter))
|
|
for key, cond := range filter {
|
|
if strings.HasPrefix(key, "$") {
|
|
sub, ok := cond.([]interface{})
|
|
if !ok {
|
|
return "", fmt.Errorf("%s expects a list", key)
|
|
}
|
|
exprs := make([]string, 0, len(sub))
|
|
for _, item := range sub {
|
|
m, ok := item.(map[string]interface{})
|
|
if !ok {
|
|
return "", fmt.Errorf("%s expects objects", key)
|
|
}
|
|
e, err := filterExpr(m)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
exprs = append(exprs, e)
|
|
}
|
|
joiner := " AND "
|
|
if key == "$or" {
|
|
joiner = " OR "
|
|
}
|
|
if key != "$and" && key != "$or" {
|
|
return "", fmt.Errorf("unsupported operator %s", key)
|
|
}
|
|
parts = append(parts, "("+strings.Join(exprs, joiner)+")")
|
|
continue
|
|
}
|
|
clause, err := fieldClause(key, cond)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
parts = append(parts, clause)
|
|
}
|
|
return strings.Join(parts, " AND "), nil
|
|
}
|
|
|
|
// fieldAccess renders the operand's comparison column over the serialized
|
|
// json_extract text. Strings and booleans compare the serialized form
|
|
// directly, which is exact-typed and never errors — booleans serialize as
|
|
// exactly 'true'/'false', and a cast would also accept the '1'/'0' a stored
|
|
// number produces. Numbers go through TRY_CAST, which is NULL on any other
|
|
// stored type. The json_get_* extractors are not used — they convert
|
|
// numeric-looking strings and booleans, and error on mismatched types.
|
|
func fieldAccess(field string, operand interface{}) (string, error) {
|
|
path := sqlEscape(existsPath(field))
|
|
switch operand.(type) {
|
|
case float64:
|
|
return fmt.Sprintf("TRY_CAST(json_extract(`metadata`, '%s') AS DOUBLE)", path), nil
|
|
case bool, string:
|
|
return fmt.Sprintf("json_extract(`metadata`, '%s')", path), nil
|
|
default:
|
|
return "", fmt.Errorf("unsupported filter operand %v", operand)
|
|
}
|
|
}
|
|
|
|
// fieldMatch renders "field <rhs>".
|
|
func fieldMatch(field string, operand interface{}, rhs string) (string, error) {
|
|
col, err := fieldAccess(field, operand)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return col + " " + rhs, nil
|
|
}
|
|
|
|
// existsPath renders the field as a JSONPath for the json_exists and
|
|
// json_extract parsers: bracket notation keeps dots, spaces, and other
|
|
// specials inside the key literal.
|
|
func existsPath(field string) string {
|
|
field = strings.ReplaceAll(field, `\`, `\\`)
|
|
field = strings.ReplaceAll(field, `"`, `\"`)
|
|
return `$["` + field + `"]`
|
|
}
|
|
|
|
func sqlEscape(s string) string {
|
|
return strings.ReplaceAll(s, "'", "''")
|
|
}
|
|
|
|
func fieldClause(field string, cond interface{}) (string, error) {
|
|
switch c := cond.(type) {
|
|
case map[string]interface{}:
|
|
parts := []string{}
|
|
for op, operand := range c {
|
|
switch op {
|
|
case "$exists":
|
|
if b, ok := operand.(bool); ok {
|
|
expr := fmt.Sprintf("json_exists(`metadata`, '%s')", sqlEscape(existsPath(field)))
|
|
if !b {
|
|
expr = "NOT " + expr
|
|
}
|
|
parts = append(parts, expr)
|
|
} else {
|
|
return "", fmt.Errorf("$exists expects a boolean")
|
|
}
|
|
case "$eq", "$ne":
|
|
lit, err := sqlLiteral(operand)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
match, err := fieldMatch(field, operand, "= "+lit)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if op == "$ne" {
|
|
match = "NOT COALESCE(" + match + ", false)"
|
|
}
|
|
parts = append(parts, match)
|
|
case "$in", "$nin":
|
|
list, ok := operand.([]interface{})
|
|
if !ok {
|
|
return "", fmt.Errorf("%s expects a list", op)
|
|
}
|
|
if len(list) == 0 {
|
|
if op == "$in" {
|
|
parts = append(parts, "false")
|
|
} else {
|
|
parts = append(parts, "true")
|
|
}
|
|
continue
|
|
}
|
|
// Members may mix types; each extractor only sees its own. A
|
|
// member whose typed extractor returns NULL must not let NOT
|
|
// invert NULL into NULL under $nin, so every group coalesces.
|
|
var groups []string
|
|
for _, item := range list {
|
|
lit, err := sqlLiteral(item)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
g, err := fieldMatch(field, item, "in ("+lit+")")
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
groups = append(groups, fmt.Sprintf("COALESCE(%s, false)", g))
|
|
}
|
|
expr := "(" + strings.Join(groups, " OR ") + ")"
|
|
if op == "$nin" {
|
|
expr = "NOT " + expr
|
|
}
|
|
parts = append(parts, expr)
|
|
case "$lt", "$lte", "$gt", "$gte":
|
|
if _, ok := operand.(float64); !ok {
|
|
return "", fmt.Errorf("%s expects a number", op)
|
|
}
|
|
lit, err := sqlLiteral(operand)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
sym := map[string]string{"$lt": "<", "$lte": "<=", "$gt": ">", "$gte": ">="}[op]
|
|
col, err := fieldAccess(field, operand)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
parts = append(parts, fmt.Sprintf("%s %s %s", col, sym, lit))
|
|
default:
|
|
return "", fmt.Errorf("unsupported operator %s", op)
|
|
}
|
|
}
|
|
return strings.Join(parts, " AND "), nil
|
|
default:
|
|
lit, err := sqlLiteral(cond)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return fieldMatch(field, cond, "= "+lit)
|
|
}
|
|
}
|
|
|
|
// jsonbString renders s as the JSON text json_extract returns for a stored
|
|
// string: quoted, with the same escapes the jsonb serializer emits.
|
|
func jsonbString(s string) string {
|
|
var b strings.Builder
|
|
b.WriteByte('"')
|
|
for _, r := range s {
|
|
switch r {
|
|
case '"':
|
|
b.WriteString(`\"`)
|
|
case '\\':
|
|
b.WriteString(`\\`)
|
|
case '\n':
|
|
b.WriteString(`\n`)
|
|
case '\r':
|
|
b.WriteString(`\r`)
|
|
case '\t':
|
|
b.WriteString(`\t`)
|
|
case 0:
|
|
b.WriteString(`\0`)
|
|
default:
|
|
if r < 0x20 || (r >= 0x7f && r <= 0x9f) {
|
|
fmt.Fprintf(&b, `\u{%x}`, r)
|
|
} else {
|
|
b.WriteRune(r)
|
|
}
|
|
}
|
|
}
|
|
b.WriteByte('"')
|
|
return b.String()
|
|
}
|
|
|
|
func sqlLiteral(v interface{}) (string, error) {
|
|
switch t := v.(type) {
|
|
case string:
|
|
return "'" + sqlEscape(jsonbString(t)) + "'", nil
|
|
case bool:
|
|
if t {
|
|
return "'true'", nil
|
|
}
|
|
return "'false'", nil
|
|
case float64:
|
|
return fmt.Sprintf("%v", t), nil
|
|
}
|
|
return "", fmt.Errorf("unsupported filter operand %v", v)
|
|
}
|
|
|
|
func (s *Server) handleQueryVectors(w http.ResponseWriter, r *http.Request) {
|
|
var req QueryVectorsRequest
|
|
if err := readJSON(r, &req); err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
arns, err := resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
var queryVector []float32
|
|
if req.QueryVector != nil {
|
|
queryVector = req.QueryVector.Float32
|
|
}
|
|
if len(queryVector) == 0 {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "queryVector is required")
|
|
return
|
|
}
|
|
record, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index)
|
|
if err != nil {
|
|
catalogErr(w, err)
|
|
return
|
|
}
|
|
idx := indexFromRecord(s, arns.bucket, arns.index, record)
|
|
if key, blocked := nonFilterableField(idx.MetadataConfiguration, req.Filter); blocked {
|
|
writeError(w, http.StatusBadRequest, "ValidationException",
|
|
fmt.Sprintf("metadata key %q is not filterable", key))
|
|
return
|
|
}
|
|
metric := idx.DistanceMetric
|
|
filter, err := translateFilter(req.Filter)
|
|
if err != nil {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", err.Error())
|
|
return
|
|
}
|
|
if req.TopK <= 0 {
|
|
writeError(w, http.StatusBadRequest, "ValidationException", "topK must be positive")
|
|
return
|
|
}
|
|
resp, err := s.vectors.QueryVectors(r.Context(), &vector_pb.QueryVectorsRequest{
|
|
Table: tableID(arns.bucket, arns.index),
|
|
QueryVector: queryVector,
|
|
TopK: uint32(req.TopK),
|
|
Filter: filter,
|
|
ReturnMetadata: req.ReturnMetadata,
|
|
ReturnDistance: req.ReturnDistance,
|
|
})
|
|
if err != nil {
|
|
vectorErr(w, err)
|
|
return
|
|
}
|
|
out := &QueryVectorsResponse{DistanceMetric: metric}
|
|
for _, v := range resp.Vectors {
|
|
o := &QueryOutputVector{Key: v.Key}
|
|
if req.ReturnDistance {
|
|
o.Distance = &v.Distance
|
|
}
|
|
if req.ReturnMetadata {
|
|
o.Metadata = structToMap(v.Metadata)
|
|
}
|
|
out.Vectors = append(out.Vectors, o)
|
|
}
|
|
writeJSON(w, http.StatusOK, out)
|
|
}
|