diff --git a/test/s3vectors/s3vectors_test.py b/test/s3vectors/s3vectors_test.py index e8138e7e35..53ce40c4dd 100644 --- a/test/s3vectors/s3vectors_test.py +++ b/test/s3vectors/s3vectors_test.py @@ -58,7 +58,7 @@ def test_vector_lifecycle(vectors): assert [h["key"] for h in hits] == ["a", "c"] distances = {h["key"]: h["distance"] for h in hits} assert distances["a"] == pytest.approx(0.0, abs=1e-6) - assert distances["c"] == pytest.approx(0.02, abs=1e-4) + assert distances["c"] == pytest.approx(0.02 ** 0.5, abs=1e-4) assert all(h["metadata"]["color"] == "red" for h in hits) vectors.delete_vectors(vectorBucketName=BUCKET, indexName=INDEX, keys=["b"]) @@ -223,6 +223,13 @@ def test_filter_arrays(vectors, index): assert matches({"$nin": [1, "2.5"]}) == ["flag", "scalar", "space"] +def test_euclidean_distance(vectors, index): + vectors.put_vectors(**index, vectors=[{"key": "p", "data": {"float32": [3, 4, 0, 0]}}]) + result = vectors.query_vectors(**index, queryVector={"float32": [0, 0, 0, 0]}, topK=1, returnDistance=True) + assert result["distanceMetric"] == "euclidean" + assert result["vectors"][0]["distance"] == pytest.approx(5.0) + + def test_cosine_metric(vectors): bucket = "cosine-" + uuid.uuid4().hex[:16] vectors.create_vector_bucket(vectorBucketName=bucket) diff --git a/weed/s3api/s3vectors/distance_metric_test.go b/weed/s3api/s3vectors/distance_metric_test.go index e3a149778e..502216bb8c 100644 --- a/weed/s3api/s3vectors/distance_metric_test.go +++ b/weed/s3api/s3vectors/distance_metric_test.go @@ -3,6 +3,7 @@ package s3vectors import ( "context" "encoding/json" + "fmt" "net/http" "testing" @@ -86,3 +87,54 @@ func TestIndexFromRecordLowercasesMetric(t *testing.T) { } } } + +type distancesFake struct { + *fakeVectors + distances []float32 +} + +func (f distancesFake) QueryVectors(ctx context.Context, req *vector_pb.QueryVectorsRequest, opts ...grpc.CallOption) (*vector_pb.QueryVectorsResponse, error) { + resp := &vector_pb.QueryVectorsResponse{} + for i, d := range f.distances { + resp.Vectors = append(resp.Vectors, &vector_pb.QueryResult{Key: fmt.Sprint(i), Distance: d}) + } + return resp, nil +} + +// The engine returns Lance's squared L2 distance; a euclidean index reports +// its square root, the other metrics their engine value. +func TestQueryVectorsReportsEuclideanDistance(t *testing.T) { + for metric, want := range map[string][]float32{ + "euclidean": {5, 0.5, 0}, + "cosine": {25, 0.25, -1e-7}, + "dotproduct": {25, 0.25, -1e-7}, + } { + fs := s3tablestest.Start(t) + engine := distancesFake{newFakeVectors(0), []float32{25, 0.25, -1e-7}} + 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: 2, DistanceMetric: metric, + }); w.Code != http.StatusOK { + t.Fatalf("create index: %d %s", w.Code, w.Body.String()) + } + w := postJSON(t, s.handleQueryVectors, &QueryVectorsRequest{ + VectorBucketName: "vec-bucket", IndexName: "idx", + QueryVector: &QueryVector{Float32: []float32{1, 0}}, TopK: 3, ReturnDistance: true, + }) + var query QueryVectorsResponse + if err := json.Unmarshal(w.Body.Bytes(), &query); err != nil || len(query.Vectors) != len(want) { + t.Fatalf("%s QueryVectors: %d %s", metric, w.Code, w.Body.String()) + } + for i, v := range query.Vectors { + if v.Distance == nil { + t.Fatalf("%s hit %d has no distance", metric, i) + } + if *v.Distance != want[i] { + t.Fatalf("%s hit %d distance %v, want %v", metric, i, *v.Distance, want[i]) + } + } + } +} diff --git a/weed/s3api/s3vectors/handlers.go b/weed/s3api/s3vectors/handlers.go index 2dbe75eab3..593e4237ce 100644 --- a/weed/s3api/s3vectors/handlers.go +++ b/weed/s3api/s3vectors/handlers.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "math" "net/http" "strconv" "strings" @@ -1938,7 +1939,12 @@ func (s *Server) handleQueryVectors(w http.ResponseWriter, r *http.Request) { for _, v := range resp.Vectors { o := &QueryOutputVector{Key: v.Key} if req.ReturnDistance { - o.Distance = &v.Distance + d := v.Distance + if metric == "euclidean" { + // Lance's L2 distance is the square of the Euclidean distance. + d = float32(math.Sqrt(math.Max(0, float64(d)))) + } + o.Distance = &d } if req.ReturnMetadata { o.Metadata = structToMap(v.Metadata)