From 8e44178dc6a2c048b0475835b49d48311cb419d5 Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Thu, 29 Oct 2026 11:14:30 -0700 Subject: [PATCH] s3vectors: cap PutVectors at 500 items and check data against the index dimension (#12200) * s3vectors: cap PutVectors at 500 items and check data against the index dimension A batch could carry any number of vectors and a vector of the wrong dimensionality reached the engine: AWS caps a PutVectors call at 500 items and requires each vector to match the index's declared dimension. * s3vectors: cover the PutVectors batch and dimension checks * s3vectors: reject query vectors that do not match the index dimension Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * s3vectors: cover the PutVectors batch size bounds Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> * s3vectors: cover the catalog-backed PutVectors handler path Run PutVectors through the real handlers over the in-memory catalog: a matching vector reaches the engine, a mismatching one is rejected with ValidationException before PutVectors is called, and a full vectorsMaxPutBatch batch is accepted. --------- Co-authored-by: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- weed/s3api/s3vectors/handlers.go | 52 ++++++--- weed/s3api/s3vectors/handlers_test.go | 149 +++++++++++++++++++++++++- 2 files changed, 184 insertions(+), 17 deletions(-) diff --git a/weed/s3api/s3vectors/handlers.go b/weed/s3api/s3vectors/handlers.go index 6e57250ca1..8a627a486a 100644 --- a/weed/s3api/s3vectors/handlers.go +++ b/weed/s3api/s3vectors/handlers.go @@ -723,7 +723,13 @@ func (s *Server) handlePutVectors(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) return } - if _, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index); err != nil { + if len(req.Vectors) == 0 || len(req.Vectors) > vectorsMaxPutBatch { + writeError(w, http.StatusBadRequest, "ValidationException", + fmt.Sprintf("vectors must contain between 1 and %d items", vectorsMaxPutBatch)) + return + } + record, err := s.indexRecord(r.Context(), r, arns.bucket, arns.index) + if err != nil { catalogErr(w, err) return } @@ -733,26 +739,15 @@ func (s *Server) handlePutVectors(w http.ResponseWriter, r *http.Request) { catalogErr(w, err) return } + dimension := indexFromRecord(s, arns.bucket, arns.index, record).Dimension 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) + pbVector, err := toPutVector(v, dimension) 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}) + vectors = append(vectors, pbVector) } _, err = s.vectors.PutVectors(r.Context(), &vector_pb.PutVectorsRequest{ Table: tableID(arns.bucket, arns.index), @@ -765,6 +760,27 @@ func (s *Server) handlePutVectors(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusOK, &PutVectorsResponse{}) } +func toPutVector(v *PutInputVector, dimension int32) (*vector_pb.Vector, error) { + if v == nil { + return nil, fmt.Errorf("vectors contains a null element") + } + if v.Key == "" { + return nil, fmt.Errorf("vector key must not be empty") + } + metadata, err := metadataStruct(v.Metadata) + if err != nil { + return nil, err + } + var data []float32 + if v.Data != nil { + data = v.Data.Float32 + } + if int32(len(data)) != dimension { + return nil, fmt.Errorf("vector %q has %d dimensions, the index declares %d", v.Key, len(data), dimension) + } + return &vector_pb.Vector{Key: v.Key, Data: data, Metadata: metadata}, nil +} + func toOutputVectors(vectors []*vector_pb.Vector, returnData, returnMetadata bool) []*OutputVector { out := make([]*OutputVector, 0, len(vectors)) for _, v := range vectors { @@ -880,6 +896,7 @@ func (s *Server) handleListVectors(w http.ResponseWriter, r *http.Request) { // one does. const ( vectorsDefaultPage = 500 + vectorsMaxPutBatch = 500 vectorsMaxPage = 1000 ) @@ -1246,6 +1263,11 @@ func (s *Server) handleQueryVectors(w http.ResponseWriter, r *http.Request) { return } idx := indexFromRecord(s, arns.bucket, arns.index, record) + if int32(len(queryVector)) != idx.Dimension { + writeError(w, http.StatusBadRequest, "ValidationException", + fmt.Sprintf("queryVector has %d dimensions, the index declares %d", len(queryVector), idx.Dimension)) + return + } if key, blocked := nonFilterableField(idx.MetadataConfiguration, req.Filter); blocked { writeError(w, http.StatusBadRequest, "ValidationException", fmt.Sprintf("metadata key %q is not filterable", key)) diff --git a/weed/s3api/s3vectors/handlers_test.go b/weed/s3api/s3vectors/handlers_test.go index f9b2433805..542327e66b 100644 --- a/weed/s3api/s3vectors/handlers_test.go +++ b/weed/s3api/s3vectors/handlers_test.go @@ -1,19 +1,30 @@ package s3vectors import ( + "bytes" "context" + "encoding/json" "fmt" + "net/http" + "net/http/httptest" "strconv" + "strings" "testing" "github.com/seaweedfs/seaweedfs/weed/pb/vector_pb" + "github.com/seaweedfs/seaweedfs/weed/s3api/s3err" + "github.com/seaweedfs/seaweedfs/weed/s3api/s3tables" + "github.com/seaweedfs/seaweedfs/weed/s3api/s3tables/s3tablestest" "google.golang.org/grpc" ) -// fakeVectors serves a fixed row set with the engine's offset-token paging. +// fakeVectors serves a fixed row set with the engine's offset-token paging +// and records the mutating calls a handler test asserts on. type fakeVectors struct { vector_pb.VectorServiceClient - rows []*vector_pb.Vector + rows []*vector_pb.Vector + created []*vector_pb.CreateIndexRequest + puts []*vector_pb.PutVectorsRequest } func newFakeVectors(n int) *fakeVectors { @@ -54,6 +65,16 @@ func (f *fakeVectors) ListVectors(ctx context.Context, req *vector_pb.ListVector return resp, nil } +func (f *fakeVectors) CreateIndex(ctx context.Context, req *vector_pb.CreateIndexRequest, opts ...grpc.CallOption) (*vector_pb.CreateIndexResponse, error) { + f.created = append(f.created, req) + return &vector_pb.CreateIndexResponse{}, nil +} + +func (f *fakeVectors) PutVectors(ctx context.Context, req *vector_pb.PutVectorsRequest, opts ...grpc.CallOption) (*vector_pb.PutVectorsResponse, error) { + f.puts = append(f.puts, req) + return &vector_pb.PutVectorsResponse{}, nil +} + // Following every segment's tokens must visit each vector exactly once. func TestListSegmentVectorsCoversEveryRowOnce(t *testing.T) { f := newFakeVectors(2500) @@ -344,3 +365,127 @@ func TestFilterFieldsNestedOperators(t *testing.T) { t.Fatal("an empty blocked key should still reject") } } + +func TestToPutVectorDimensionCheck(t *testing.T) { + good := &PutInputVector{Key: "a", Data: &VectorData{Float32: []float32{1, 2, 3}}} + for _, tc := range []struct { + name string + v *PutInputVector + dim int32 + ok bool + }{ + {"matching", good, 3, true}, + {"mismatched", good, 4, false}, + {"missing data", &PutInputVector{Key: "a"}, 3, false}, + {"empty data", &PutInputVector{Key: "a", Data: &VectorData{}}, 3, false}, + {"null element", nil, 3, false}, + {"empty key", &PutInputVector{Data: &VectorData{Float32: []float32{1, 2, 3}}}, 3, false}, + } { + t.Run(tc.name, func(t *testing.T) { + if _, err := toPutVector(tc.v, tc.dim); (err == nil) != tc.ok { + t.Fatalf("toPutVector err = %v, want ok %v", err, tc.ok) + } + }) + } +} + +func TestPutVectorsRejectsBadBatchSizes(t *testing.T) { + s := &Server{} + for _, n := range []int{0, vectorsMaxPutBatch + 1} { + vectors := make([]*PutInputVector, n) + for i := range vectors { + vectors[i] = &PutInputVector{} + } + body, _ := json.Marshal(&PutVectorsRequest{ + VectorBucketName: "bucket", + IndexName: "index", + Vectors: vectors, + }) + w := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodPost, "/PutVectors", bytes.NewReader(body)) + s.handlePutVectors(w, r) + if w.Code != http.StatusBadRequest { + t.Fatalf("batch of %d: status %d, want %d", n, w.Code, http.StatusBadRequest) + } + } +} + +// allowAllAuth stands in for the zero-config deployment NewServer wires +// DefaultAllow from. +type allowAllAuth struct{} + +func (allowAllAuth) AuthenticateRequest(r *http.Request) (string, interface{}, s3err.ErrorCode) { + return "", nil, s3err.ErrNone +} +func (allowAllAuth) DefaultAllow() bool { return true } + +func postJSON(t *testing.T, handler http.HandlerFunc, payload interface{}) *httptest.ResponseRecorder { + t.Helper() + body, err := json.Marshal(payload) + if err != nil { + t.Fatalf("marshal request: %v", err) + } + w := httptest.NewRecorder() + handler(w, httptest.NewRequest(http.MethodPost, "/", bytes.NewReader(body))) + return w +} + +// The handler must take the declared dimension from the catalog record: a +// matching vector reaches the engine, a mismatching one is rejected without a +// PutVectors call, and a full vectorsMaxPutBatch batch is accepted. +func TestPutVectorsUsesCatalogDimension(t *testing.T) { + fs := s3tablestest.Start(t) + engine := &fakeVectors{} + 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()) + } + if len(engine.created) != 1 || engine.created[0].Dimension != 3 { + t.Fatalf("engine CreateIndex calls = %v, want one dimension-3 index", engine.created) + } + + put := func(vectors ...*PutInputVector) *httptest.ResponseRecorder { + return postJSON(t, s.handlePutVectors, &PutVectorsRequest{ + VectorBucketName: "vec-bucket", + IndexName: "idx", + Vectors: vectors, + }) + } + + matching := &PutInputVector{Key: "a", Data: &VectorData{Float32: []float32{1, 2, 3}}} + if w := put(matching); w.Code != http.StatusOK { + t.Fatalf("matching vector: %d %s", w.Code, w.Body.String()) + } + if len(engine.puts) != 1 || len(engine.puts[0].Vectors) != 1 { + t.Fatalf("engine PutVectors calls = %v, want one call of one vector", engine.puts) + } + + mismatched := &PutInputVector{Key: "b", Data: &VectorData{Float32: []float32{1, 2, 3, 4}}} + if w := put(mismatched); w.Code != http.StatusBadRequest || !strings.Contains(w.Body.String(), "ValidationException") { + t.Fatalf("mismatched vector: %d %s, want 400 ValidationException", w.Code, w.Body.String()) + } + if len(engine.puts) != 1 { + t.Fatalf("rejected vector reached the engine: %d calls", len(engine.puts)) + } + + batch := make([]*PutInputVector, vectorsMaxPutBatch) + for i := range batch { + batch[i] = &PutInputVector{Key: fmt.Sprintf("k%d", i), Data: &VectorData{Float32: []float32{1, 2, 3}}} + } + if w := put(batch...); w.Code != http.StatusOK { + t.Fatalf("%d-vector batch: %d %s", vectorsMaxPutBatch, w.Code, w.Body.String()) + } + if len(engine.puts) != 2 || len(engine.puts[1].Vectors) != vectorsMaxPutBatch { + t.Fatalf("engine PutVectors calls = %d, want a second call of %d vectors", len(engine.puts), vectorsMaxPutBatch) + } +}