Merge branch 'develop' into combined-bot-prs-branch

Author: shahzadlone

internal/index/hnsw/heap.go

@@ -18,7 +18,7 @@ import "container/heap"
 // by the neighbour-selection heuristic).
 type candidate struct {
 	id     NodeID
-	dist   float32
+	dist   float64
 	vector []float32
 }
 

internal/index/hnsw/heap_test.go

@@ -26,28 +26,28 @@ type heapUnderTest struct {
 	// new returns a fresh, empty heap.Interface backed by the concrete heap type.
 	new func() heap.Interface
 	// root reads h[0].dist from the concrete heap value.
-	root func(h heap.Interface) float32
+	root func(h heap.Interface) float64
 	// popOrder returns the distances in the order Pop yields them: ascending for minHeap,
 	// descending for maxHeap.
-	popOrder func(sorted []float32) []float32
+	popOrder func(sorted []float64) []float64
 	// extreme returns the element this heap would pop next from the given contents: the minimum
 	// for minHeap, the maximum for maxHeap. Panics on empty input (never called that way).
-	extreme func(contents []float32) float32
+	extreme func(contents []float64) float64
 }
 
 func minHeapUT() heapUnderTest {
 	return heapUnderTest{
 		name: "minHeap",
 		new:  func() heap.Interface { return newMinHeap() },
-		root: func(h heap.Interface) float32 {
+		root: func(h heap.Interface) float64 {
 			mh, _ := h.(*candidateHeap)
 			return mh.items[0].dist
 		},
-		popOrder: func(sorted []float32) []float32 {
-			out := append([]float32(nil), sorted...) // sorted is ascending → min pops ascending
+		popOrder: func(sorted []float64) []float64 {
+			out := append([]float64(nil), sorted...) // sorted is ascending → min pops ascending
 			return out
 		},
-		extreme: func(contents []float32) float32 {
+		extreme: func(contents []float64) float64 {
 			m := contents[0]
 			for _, v := range contents {
 				if v < m {
@@ -63,18 +63,18 @@ func maxHeapUT() heapUnderTest {
 	return heapUnderTest{
 		name: "maxHeap",
 		new:  func() heap.Interface { return newMaxHeap() },
-		root: func(h heap.Interface) float32 {
+		root: func(h heap.Interface) float64 {
 			mh, _ := h.(*candidateHeap)
 			return mh.items[0].dist
 		},
-		popOrder: func(sorted []float32) []float32 {
-			out := make([]float32, len(sorted)) // sorted is ascending → max pops descending
+		popOrder: func(sorted []float64) []float64 {
+			out := make([]float64, len(sorted)) // sorted is ascending → max pops descending
 			for i, v := range sorted {
 				out[len(sorted)-1-i] = v
 			}
 			return out
 		},
-		extreme: func(contents []float32) float32 {
+		extreme: func(contents []float64) float64 {
 			m := contents[0]
 			for _, v := range contents {
 				if v > m {
@@ -94,9 +94,9 @@ func popCandidate(t *testing.T, h heap.Interface) candidate {
 	return c
 }
 
-func popAllDists(t *testing.T, h heap.Interface) []float32 {
+func popAllDists(t *testing.T, h heap.Interface) []float64 {
 	t.Helper()
-	out := make([]float32, 0, h.Len())
+	out := make([]float64, 0, h.Len())
 	for h.Len() > 0 {
 		out = append(out, popCandidate(t, h).dist)
 	}
@@ -106,7 +106,7 @@ func popAllDists(t *testing.T, h heap.Interface) []float32 {
 func TestHeap_PushThenPopAll_YieldsOrderedDistances(t *testing.T) {
 	// A representative unordered input; sorted ascending is the reference the per-heap popOrder
 	// derives its expectation from.
-	inputs := map[string][]float32{
+	inputs := map[string][]float64{
 		"unordered distinct":  {5, 1, 4, 2, 3},
 		"already ascending":   {1, 2, 3, 4},
 		"already descending":  {4, 3, 2, 1},
@@ -134,7 +134,7 @@ func TestHeap_PushThenPopAll_YieldsOrderedDistances(t *testing.T) {
 }
 
 func TestHeap_Root_IsExtremeElement(t *testing.T) {
-	in := []float32{5, 1, 4, 2, 3}
+	in := []float64{5, 1, 4, 2, 3}
 
 	for _, ut := range []heapUnderTest{minHeapUT(), maxHeapUT()} {
 		t.Run(ut.name, func(t *testing.T) {
@@ -157,9 +157,9 @@ func TestHeap_InterleavedPushPop_ReturnsCurrentExtreme(t *testing.T) {
 	for _, ut := range []heapUnderTest{minHeapUT(), maxHeapUT()} {
 		t.Run(ut.name, func(t *testing.T) {
 			h := ut.new()
-			model := []float32{}
+			model := []float64{}
 
-			push := func(d float32) {
+			push := func(d float64) {
 				heap.Push(h, candidate{dist: d})
 				model = append(model, d)
 			}
@@ -194,7 +194,7 @@ func TestHeap_PreservesCandidatePayload(t *testing.T) {
 }
 
 // removeFirst returns s with the first occurrence of v removed (reference-model bookkeeping).
-func removeFirst(s []float32, v float32) []float32 {
+func removeFirst(s []float64, v float64) []float64 {
 	for i, x := range s {
 		if x == v {
 			return append(s[:i:i], s[i+1:]...)
@@ -204,8 +204,8 @@ func removeFirst(s []float32, v float32) []float32 {
 }
 
 // ascendingSorted returns a new ascending-sorted copy of in, used as the reference ordering.
-func ascendingSorted(in []float32) []float32 {
-	out := append([]float32(nil), in...)
+func ascendingSorted(in []float64) []float64 {
+	out := append([]float64(nil), in...)
 	for i := 1; i < len(out); i++ {
 		for j := i; j > 0 && out[j] < out[j-1]; j-- {
 			out[j], out[j-1] = out[j-1], out[j]

internal/index/hnsw/hnsw.go

@@ -136,7 +136,11 @@ func dot(a, b []float32) float64 {
 // distance returns the distance between two vectors under the given metric, smaller being nearer.
 // Both must have come through vectorForMetric for this metric. An unrecognised metric is a
 // programming error rather than a silent fallback, so it panics.
-func distance(metric Metric, a, b []float32) float32 {
+//
+// The result is float64 even though vectors are stored as float32. Squaring a float32 can exceed
+// what a float32 holds, and every such distance would land on +Inf and compare equal, leaving the
+// documents in an arbitrary order. Float64 has the range to keep them apart.
+func distance(metric Metric, a, b []float32) float64 {
 	switch metric {
 	case Cosine:
 		return cosineDistance(a, b)
@@ -150,28 +154,25 @@ func distance(metric Metric, a, b []float32) float32 {
 }
 
 // cosineDistance computes 1 - dot(a, b) for unit-length vectors a and b.
-func cosineDistance(a, b []float32) float32 {
-	return float32(1 - dot(a, b))
+func cosineDistance(a, b []float32) float64 {
+	return 1 - dot(a, b)
 }
 
 // squaredEuclideanDistance computes the squared straight-line distance between a and b. The square
 // root is skipped: it is monotonic, and distances are only ever compared against each other. If the
 // lengths differ, only the shared leading elements are used, matching dot.
-//
-// Squaring costs range: the sum is held in float64, but the result narrows to float32, so components
-// beyond roughly 1e19 saturate to +Inf. Such a vector still sorts as the farthest, so ordering holds.
-func squaredEuclideanDistance(a, b []float32) float32 {
+func squaredEuclideanDistance(a, b []float32) float64 {
 	var sum float64
 	n := min(len(a), len(b))
 	for i := range n {
 		d := float64(a[i]) - float64(b[i])
 		sum += d * d
 	}
-	return float32(sum)
+	return sum
 }
 
 // dotProductDistance computes -dot(a, b). The dot product is a similarity (bigger is nearer), so the
 // sign is flipped to keep smaller nearer.
-func dotProductDistance(a, b []float32) float32 {
-	return float32(-dot(a, b))
+func dotProductDistance(a, b []float32) float64 {
+	return -dot(a, b)
 }

internal/index/hnsw/hnsw_test.go

@@ -26,7 +26,7 @@ import (
 func bruteForceKNN(metric Metric, query []float32, vectors map[NodeID][]float32, k int) []NodeID {
 	type scored struct {
 		id   NodeID
-		dist float32
+		dist float64
 	}
 	scoredList := make([]scored, 0, len(vectors))
 	for id, v := range vectors {
@@ -399,3 +399,38 @@ func TestGraph_DotProduct_RanksLongerVectorNearer(t *testing.T) {
 func TestDistance_UnknownMetric_Panics(t *testing.T) {
 	assert.Panics(t, func() { distance(Metric(99), []float32{1}, []float32{1}) })
 }
+
+// Squaring a large float32 exceeds what a float32 holds. Held in one, every distance here would be
+// +Inf and compare equal, leaving the order to chance. The query sits next to the largest vector, so
+// that one has to come first and the rest follow by how far they are.
+// https://github.com/sourcenetwork/defradb/issues/5220
+func TestGraph_Euclidean_LargeVectors_OrderByDistanceNotOverflow(t *testing.T) {
+	g := New(NewMemStore(), Euclidean, DefaultParams(8), 1)
+
+	require.NoError(t, g.Insert(1, []float32{-3.4028235e+38}))
+	require.NoError(t, g.Insert(2, []float32{0}))
+	require.NoError(t, g.Insert(3, []float32{90}))
+
+	result, err := g.Search([]float32{-3.4028235e+38}, 3, 64)
+	require.NoError(t, err)
+	require.Equal(t, []NodeID{1, 2, 3}, result)
+}
+
+// The distances themselves must stay finite, not just come out in the right order. Held in a float32
+// these all collapse onto +Inf, which reports every document as equally near and is what left the
+// order to chance.
+func TestGraph_Euclidean_LargeVectors_DistancesStayFinite(t *testing.T) {
+	g := New(NewMemStore(), Euclidean, DefaultParams(8), 1)
+
+	require.NoError(t, g.Insert(1, []float32{0}))
+	require.NoError(t, g.Insert(2, []float32{3.4028235e+38}))
+
+	neighbors, err := g.SearchWithDistance([]float32{-3.4028235e+38}, 2, 64)
+	require.NoError(t, err)
+	require.Len(t, neighbors, 2)
+	for _, n := range neighbors {
+		require.False(t, math.IsInf(n.Distance, 1), "distance overflowed to +Inf")
+	}
+	// The far vector is twice the distance away, so squared it is four times as far.
+	require.Less(t, neighbors[0].Distance, neighbors[1].Distance)
+}

internal/index/hnsw/search.go

@@ -44,7 +44,7 @@ func (g *HNSWIndex) SearchWithDistance(query []float32, k, efSearch int) ([]Neig
 	}
 	out := make([]Neighbor, len(w))
 	for i, c := range w {
-		out[i] = Neighbor{ID: c.id, Distance: float64(c.dist)}
+		out[i] = Neighbor{ID: c.id, Distance: c.dist}
 	}
 	return out, nil
 }

tests/integration/index/vector_metrics_test.go

@@ -262,3 +262,111 @@ func TestVectorIndex_DropThenRecreateWithDifferentMetric_IsAllowed(t *testing.T)
 
 	testUtils.ExecuteTestCase(t, test)
 }
+
+// Vectors whose squared distance exceeds a float32 are searchable and scored correctly. The graph's
+// own ordering is covered by the engine tests; this covers the query path end to end, where the
+// score is computed separately and was never affected.
+// https://github.com/sourcenetwork/defradb/issues/5220
+func TestVectorIndex_EuclideanOnLargeVectors_OrdersByDistance(t *testing.T) {
+	test := testUtils.TestCase{
+		Actions: []any{
+			&action.AddCollection{
+				SDL: `type User {
+					name: String
+					vector: [Float32!] @index(vector: {dimensions: 1, hnsw: {metric: EUCLIDEAN}})
+				}`,
+			},
+			&action.AddDoc{DocMap: map[string]any{"name": "mid", "vector": []float32{2e19}}},
+			&action.AddDoc{DocMap: map[string]any{"name": "far", "vector": []float32{3e19}}},
+			&action.AddDoc{DocMap: map[string]any{"name": "farthest", "vector": []float32{4e19}}},
+			&action.WaitForIndexReady{CollectionID: 0},
+			&action.Request{
+				Request: `query {
+					User(order: {_alias: {sim: DESC}}, limit: 1){
+						name
+						sim: SIMILARITY(vector: {vector: [0]})
+					}
+				}`,
+				Results: map[string]any{
+					"User": []map[string]any{
+						{"name": "mid", "sim": -3.999999984405158e+38},
+					},
+				},
+			},
+		},
+	}
+
+	testUtils.ExecuteTestCase(t, test)
+}
+
+// The dot product of two large vectors also exceeds a float32, so it needs the same float64 handling
+// as euclidean. Cosine cannot overflow because its vectors are normalised first, but it is covered
+// alongside so a future change to either metric is caught here.
+// https://github.com/sourcenetwork/defradb/issues/5220
+func TestVectorIndex_DotProductOnLargeVectors_OrdersByDistance(t *testing.T) {
+	test := testUtils.TestCase{
+		Actions: []any{
+			&action.AddCollection{
+				SDL: `type User {
+					name: String
+					vector: [Float32!] @index(vector: {dimensions: 1, hnsw: {metric: DOT}})
+				}`,
+			},
+			&action.AddDoc{DocMap: map[string]any{"name": "near", "vector": []float32{4e19}}},
+			&action.AddDoc{DocMap: map[string]any{"name": "mid", "vector": []float32{3e19}}},
+			&action.AddDoc{DocMap: map[string]any{"name": "far", "vector": []float32{2e19}}},
+			&action.WaitForIndexReady{CollectionID: 0},
+			&action.Request{
+				Request: `query {
+					User(order: {_alias: {sim: DESC}}, limit: 3){
+						name
+						sim: SIMILARITY(vector: {vector: [1e20]})
+					}
+				}`,
+				Results: map[string]any{
+					"User": []map[string]any{
+						{"name": "near", "sim": 4.0000000723660884e+39},
+						{"name": "mid", "sim": 3.000000164225731e+39},
+						{"name": "far", "sim": 2.0000000361830442e+39},
+					},
+				},
+			},
+		},
+	}
+
+	testUtils.ExecuteTestCase(t, test)
+}
+
+func TestVectorIndex_CosineOnLargeVectors_OrdersByDistance(t *testing.T) {
+	test := testUtils.TestCase{
+		Actions: []any{
+			&action.AddCollection{
+				SDL: `type User {
+					name: String
+					vector: [Float32!] @index(vector: {dimensions: 2, hnsw: {metric: COSINE}})
+				}`,
+			},
+			&action.AddDoc{DocMap: map[string]any{"name": "aligned", "vector": []float32{2e19, 0}}},
+			&action.AddDoc{DocMap: map[string]any{"name": "diagonal", "vector": []float32{2e19, 2e19}}},
+			&action.AddDoc{DocMap: map[string]any{"name": "orthogonal", "vector": []float32{0, 2e19}}},
+			&action.WaitForIndexReady{CollectionID: 0},
+			&action.Request{
+				Request: `query {
+					User(order: {_alias: {sim: DESC}}, limit: 3){
+						name
+						sim: SIMILARITY(vector: {vector: [1e20, 0]})
+					}
+				}`,
+				Results: map[string]any{
+					"User": []map[string]any{
+						{"name": "aligned", "sim": 1.0},
+						{"name": "diagonal", "sim": 0.7071067811865476},
+						{"name": "orthogonal", "sim": 0.0},
+					},
+				},
+			},
+		},
+	}
+
+	testUtils.ExecuteTestCase(t, test)
+}