diff --git a/test/s3vectors/s3vectors_test.py b/test/s3vectors/s3vectors_test.py index abd06d3ea4..365f74060a 100644 --- a/test/s3vectors/s3vectors_test.py +++ b/test/s3vectors/s3vectors_test.py @@ -186,6 +186,8 @@ def test_list_pagination_and_segments(vectors, index): segmented = [row for segment in range(4) for row in collect(segmentCount=4, segmentIndex=segment)] assert sorted(segmented, key=lambda r: r["key"]) == rows expect_error("ValidationException", vectors.list_vectors, **index, nextToken="invalid") + for bad in ({"maxResults": 1001}, {"segmentCount": 17, "segmentIndex": 0}, {"segmentIndex": 0}): + expect_error("ValidationException", vectors.list_vectors, **index, **bad) def test_filter_types(vectors, index): diff --git a/weed/s3api/s3vectors/distance_metric_test.go b/weed/s3api/s3vectors/distance_metric_test.go new file mode 100644 index 0000000000..e3a149778e --- /dev/null +++ b/weed/s3api/s3vectors/distance_metric_test.go @@ -0,0 +1,88 @@ +package s3vectors + +import ( + "context" + "encoding/json" + "net/http" + "testing" + + "github.com/seaweedfs/seaweedfs/weed/pb/vector_pb" + "github.com/seaweedfs/seaweedfs/weed/s3api/s3tables" + "github.com/seaweedfs/seaweedfs/weed/s3api/s3tables/s3tablestest" + "google.golang.org/grpc" +) + +type queryVectorsFake struct{ *fakeVectors } + +func (queryVectorsFake) QueryVectors(ctx context.Context, req *vector_pb.QueryVectorsRequest, opts ...grpc.CallOption) (*vector_pb.QueryVectorsResponse, error) { + return &vector_pb.QueryVectorsResponse{}, nil +} + +// A metric sent in any case is stored, and returned by GetIndex and +// QueryVectors, as the lowercase AWS enum value. +func TestCreateIndexStoresLowercaseMetric(t *testing.T) { + fs := s3tablestest.Start(t) + engine := queryVectorsFake{newFakeVectors(0)} + s := NewServer(s3tables.NewManagerClient(fs.Client), allowAllAuth{}, engine) + if w := postJSON(t, s.handleCreateVectorBucket, &CreateVectorBucketRequest{VectorBucketName: "vec-bucket"}); w.Code != http.StatusOK { + t.Fatalf("create bucket: %d %s", w.Code, w.Body.String()) + } + index := &GetIndexRequest{VectorBucketName: "vec-bucket", IndexName: "idx"} + if w := postJSON(t, s.handleCreateIndex, &CreateIndexRequest{ + VectorBucketName: index.VectorBucketName, + IndexName: index.IndexName, + DataType: "float32", + Dimension: 2, + DistanceMetric: "COSINE", + }); w.Code != http.StatusOK { + t.Fatalf("create index: %d %s", w.Code, w.Body.String()) + } + if len(engine.created) != 1 || engine.created[0].DistanceMetric != vector_pb.DistanceMetric_COSINE { + t.Fatalf("engine CreateIndex calls = %v, want one cosine index", engine.created) + } + + record, err := s.indexRecord(context.Background(), index.VectorBucketName, index.IndexName) + if err != nil { + t.Fatal(err) + } + if got := record["metadata"].(map[string]interface{})["fullMetadata"].(map[string]interface{})["distanceMetric"]; got != "cosine" { + t.Fatalf("stored distanceMetric %v, want cosine", got) + } + + w := postJSON(t, s.handleGetIndex, index) + var got GetIndexResponse + if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil || got.Index == nil || got.Index.DistanceMetric != "cosine" { + t.Fatalf("GetIndex: %d %s", w.Code, w.Body.String()) + } + + w = postJSON(t, s.handleQueryVectors, &QueryVectorsRequest{ + VectorBucketName: index.VectorBucketName, + IndexName: index.IndexName, + QueryVector: &QueryVector{Float32: []float32{1, 0}}, + TopK: 1, + }) + var query QueryVectorsResponse + if err := json.Unmarshal(w.Body.Bytes(), &query); err != nil || query.DistanceMetric != "cosine" { + t.Fatalf("QueryVectors: %d %s", w.Code, w.Body.String()) + } +} + +// Records stored before the metric was normalized render in lowercase too. +func TestIndexFromRecordLowercasesMetric(t *testing.T) { + s := &Server{} + for stored, want := range map[interface{}]string{ + "COSINE": "cosine", + "Euclidean": "euclidean", + "": "euclidean", + nil: "euclidean", + } { + record := map[string]interface{}{ + "metadata": map[string]interface{}{ + "fullMetadata": map[string]interface{}{"dimension": "2", "distanceMetric": stored}, + }, + } + if got := indexFromRecord(s, "b", "i", record).DistanceMetric; got != want { + t.Fatalf("stored %v: distanceMetric %q, want %q", stored, got, want) + } + } +} diff --git a/weed/s3api/s3vectors/handlers.go b/weed/s3api/s3vectors/handlers.go index bb3c65b7c3..0c20fa228d 100644 --- a/weed/s3api/s3vectors/handlers.go +++ b/weed/s3api/s3vectors/handlers.go @@ -958,7 +958,7 @@ func (s *Server) handleCreateIndex(w http.ResponseWriter, r *http.Request) { } metadata := map[string]string{ "dimension": fmt.Sprint(req.Dimension), - "distanceMetric": req.DistanceMetric, + "distanceMetric": strings.ToLower(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(), @@ -1049,7 +1049,7 @@ func indexFromRecord(s *Server, bucket, name string, m map[string]interface{}) * var dim float64 fmt.Sscan(fmt.Sprint(fm["dimension"]), &dim) idx.Dimension = int32(dim) - idx.DistanceMetric = fmt.Sprint(fm["distanceMetric"]) + idx.DistanceMetric = strings.ToLower(fmt.Sprint(fm["distanceMetric"])) switch keys := fm["nonFilterableMetadataKeys"].(type) { case string: var list []string @@ -1360,6 +1360,10 @@ func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) return } + if err := listVectorsBounds(&req); err != nil { + writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) + return + } arns, err := s.resolveIndexARNs(req.VectorBucketName, req.VectorBucketArn, req.IndexName, req.IndexArn) if err != nil { writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) @@ -1373,20 +1377,16 @@ func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) { catalogErr(w, err) return } - if req.MaxResults < 0 { - writeError(w, http.StatusBadRequest, "ValidationException", "maxResults must be non-negative") - return + maxResults := 0 + if req.MaxResults != nil { + maxResults = *req.MaxResults } // 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 - } + // the index here and keep only the caller's share; a call scans about + // segmentCount × maxResults rows. The engine's next_token is a row offset, + // so the caller's token is the offset to resume scanning from. + if req.SegmentCount != nil { var offset uint64 if req.NextToken != "" { var err error @@ -1396,7 +1396,7 @@ func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) { } } 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) + offset, maxResults, uint64(*req.SegmentCount), uint64(*req.SegmentIndex), req.ReturnData, req.ReturnMetadata) if err != nil { vectorErr(w, err) return @@ -1409,7 +1409,7 @@ func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) { } resp, err := s.vectors.ListVectors(r.Context(), &vector_pb.ListVectorsRequest{ Table: tableID(arns.bucket, arns.index), - MaxResults: uint32(req.MaxResults), + MaxResults: uint32(maxResults), NextToken: req.NextToken, ReturnData: req.ReturnData, ReturnMetadata: req.ReturnMetadata, @@ -1425,7 +1425,8 @@ func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) { } // The page constants mirror the engine's page sizing so a segmented listing -// pages like a plain one does; vectorsMaxTopK is the AWS QueryVectors limit. +// pages like a plain one does; vectorsMaxPage is also the AWS ListVectors +// maxResults cap, and vectorsMaxTopK the AWS QueryVectors limit. const ( vectorsDefaultPage = 500 vectorsMaxGetKeys = 100 @@ -1433,8 +1434,30 @@ const ( vectorsMaxPage = 1000 vectorsMaxListPage = 500 vectorsMaxTopK = 10000 + vectorsMaxSegments = 16 ) +// listVectorsBounds enforces the AWS ListVectors windows: maxResults 1..1000, +// segmentCount 1..16, and segmentIndex below it, the two given together. +func listVectorsBounds(req *ListVectorsRequest) error { + if req.MaxResults != nil && (*req.MaxResults < 1 || *req.MaxResults > vectorsMaxPage) { + return fmt.Errorf("maxResults must be between 1 and %d", vectorsMaxPage) + } + if (req.SegmentCount == nil) != (req.SegmentIndex == nil) { + return fmt.Errorf("segmentCount and segmentIndex must be specified together") + } + if req.SegmentCount == nil { + return nil + } + if *req.SegmentCount < 1 || *req.SegmentCount > vectorsMaxSegments { + return fmt.Errorf("segmentCount must be between 1 and %d", vectorsMaxSegments) + } + if *req.SegmentIndex < 0 || *req.SegmentIndex >= *req.SegmentCount { + return fmt.Errorf("segmentIndex must be in [0, segmentCount)") + } + return nil +} + // 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 diff --git a/weed/s3api/s3vectors/list_bounds_test.go b/weed/s3api/s3vectors/list_bounds_test.go new file mode 100644 index 0000000000..cdd9e8e3b3 --- /dev/null +++ b/weed/s3api/s3vectors/list_bounds_test.go @@ -0,0 +1,101 @@ +package s3vectors + +import ( + "context" + "encoding/json" + "net/http" + "strings" + "testing" + + "github.com/seaweedfs/seaweedfs/weed/pb/vector_pb" + "github.com/seaweedfs/seaweedfs/weed/s3api/s3tables" + "github.com/seaweedfs/seaweedfs/weed/s3api/s3tables/s3tablestest" + "google.golang.org/grpc" +) + +// countingVectors counts engine ListVectors calls. +type countingVectors struct { + *fakeVectors + lists int +} + +func (c *countingVectors) ListVectors(ctx context.Context, req *vector_pb.ListVectorsRequest, opts ...grpc.CallOption) (*vector_pb.ListVectorsResponse, error) { + c.lists++ + return c.fakeVectors.ListVectors(ctx, req, opts...) +} + +func listVectorsBody(fields string) json.RawMessage { + return json.RawMessage(`{"vectorBucketName":"vec-bucket","indexName":"idx"` + fields + `}`) +} + +// Out-of-range values are rejected before authorization (the server has no +// catalog) and never reach the engine. +func TestListVectorsBoundsRejectBeforeEngine(t *testing.T) { + engine := &countingVectors{fakeVectors: newFakeVectors(50000)} + s := &Server{vectors: engine} + for name, fields := range map[string]string{ + "maxResults zero": `,"maxResults":0`, + "maxResults negative": `,"maxResults":-1`, + "maxResults over": `,"maxResults":1001`, + "segmentCount zero": `,"segmentCount":0,"segmentIndex":0`, + "segmentCount over": `,"segmentCount":17,"segmentIndex":0`, + "segmentCount huge": `,"segmentCount":1000000,"segmentIndex":999999,"maxResults":1`, + "segmentIndex negative": `,"segmentCount":4,"segmentIndex":-1`, + "segmentIndex at count": `,"segmentCount":4,"segmentIndex":4`, + "segmentIndex without count": `,"segmentIndex":0`, + "segmentCount without index": `,"segmentCount":4`, + "segment ok, maxResults over": `,"segmentCount":4,"segmentIndex":0,"maxResults":1001`, + "segment ok, maxResults zeroed": `,"segmentCount":4,"segmentIndex":0,"maxResults":0`, + } { + w := postJSON(t, s.handleListVectors, listVectorsBody(fields)) + if w.Code != http.StatusBadRequest || !strings.Contains(w.Body.String(), "ValidationException") { + t.Fatalf("%s: %d %s, want 400 ValidationException", name, w.Code, w.Body.String()) + } + } + if engine.lists != 0 { + t.Fatalf("engine saw %d ListVectors calls, want none", engine.lists) + } +} + +// In-range values list normally; a segmented page at the top of the range +// costs one engine page, not a scan of the index. +func TestListVectorsBoundsAccepted(t *testing.T) { + fs := s3tablestest.Start(t) + engine := &countingVectors{fakeVectors: newFakeVectors(50000)} + s := NewServer(s3tables.NewManagerClient(fs.Client), allowAllAuth{}, engine) + if w := postJSON(t, s.handleCreateVectorBucket, &CreateVectorBucketRequest{VectorBucketName: "vec-bucket"}); w.Code != http.StatusOK { + t.Fatalf("create bucket: %d %s", w.Code, w.Body.String()) + } + if w := postJSON(t, s.handleCreateIndex, &CreateIndexRequest{ + VectorBucketName: "vec-bucket", + IndexName: "idx", + DataType: "float32", + Dimension: 3, + DistanceMetric: "euclidean", + }); w.Code != http.StatusOK { + t.Fatalf("create index: %d %s", w.Code, w.Body.String()) + } + for _, tc := range []struct { + fields string + rows int + }{ + {"", vectorsDefaultPage}, + {`,"maxResults":1`, 1}, + {`,"maxResults":1000`, vectorsMaxPage}, + {`,"segmentCount":1,"segmentIndex":0,"maxResults":1000`, vectorsMaxPage}, + {`,"segmentCount":16,"segmentIndex":15,"maxResults":1`, 1}, + } { + engine.lists = 0 + w := postJSON(t, s.handleListVectors, listVectorsBody(tc.fields)) + if w.Code != http.StatusOK { + t.Fatalf("%q: %d %s", tc.fields, w.Code, w.Body.String()) + } + var out ListVectorsResponse + if err := json.Unmarshal(w.Body.Bytes(), &out); err != nil { + t.Fatalf("%q: decode: %v", tc.fields, err) + } + if len(out.Vectors) != tc.rows || engine.lists != 1 { + t.Fatalf("%q: %d rows in %d engine calls, want %d in 1", tc.fields, len(out.Vectors), engine.lists, tc.rows) + } + } +} diff --git a/weed/s3api/s3vectors/types.go b/weed/s3api/s3vectors/types.go index 423dc71949..a9c5b96c51 100644 --- a/weed/s3api/s3vectors/types.go +++ b/weed/s3api/s3vectors/types.go @@ -233,12 +233,12 @@ type ListVectorsRequest struct { IndexName string `json:"indexName"` VectorBucketName string `json:"vectorBucketName"` VectorBucketArn string `json:"vectorBucketArn"` - MaxResults int `json:"maxResults"` + MaxResults *int `json:"maxResults"` NextToken string `json:"nextToken"` ReturnData bool `json:"returnData"` ReturnMetadata bool `json:"returnMetadata"` - SegmentCount int `json:"segmentCount"` - SegmentIndex int `json:"segmentIndex"` + SegmentCount *int `json:"segmentCount"` + SegmentIndex *int `json:"segmentIndex"` } type ListVectorsResponse struct {