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

@@ -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))

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)
+	}
+}