fix(vector_stores/turbopuffer): make search score respect distance_metric (#6559)

This commit is contained in:
Yash Singh
2026-09-25 17:08:20 +05:30
committed by GitHub
parent 5fd01d284a
commit 8127e8bd47
2 changed files with 37 additions and 1 deletions
+8 -1
View File
@@ -122,7 +122,14 @@ class TurbopufferDB(VectorStoreBase):
dist = row_dict.pop("$dist", None)
row_dict.pop("vector", None)
score = 1 - dist if dist is not None else None
if dist is None:
score = None
elif self.distance_metric == "euclidean_squared":
# $dist is unbounded squared-L2 (lower = closer); map to a bounded
# higher-is-better score, mirroring milvus/baidu. Cosine returns 1 - dist.
score = 1.0 / (1.0 + dist)
else:
score = 1 - dist
results.append(OutputData(
id=row_id,
+29
View File
@@ -232,6 +232,35 @@ class TestParseOutput:
def test_parse_empty_rows(self, db):
assert db._parse_output([]) == []
def test_parse_cosine_score_is_one_minus_dist(self, db):
# Default cosine metric: score = 1 - dist, unchanged by the metric fix.
results = db._parse_output([_make_row("id1", dist=0.25)])
assert results[0].score == pytest.approx(0.75)
def test_parse_euclidean_squared_score_is_bounded(self, mock_client):
# euclidean_squared $dist is unbounded (e.g. 4.0). 1 - dist would give -3.0,
# violating the higher-is-better contract; map it to 1/(1+dist) instead.
db = TurbopufferDB(
collection_name="test_ns",
embedding_model_dims=4,
api_key="tpuf_test_key",
region="gcp-us-central1",
distance_metric="euclidean_squared",
)
results = db._parse_output([_make_row("id1", dist=4.0)])
assert results[0].score == pytest.approx(0.2)
assert 0.0 <= results[0].score <= 1.0
def test_parse_euclidean_squared_preserves_none(self, mock_client):
db = TurbopufferDB(
collection_name="test_ns",
embedding_model_dims=4,
api_key="tpuf_test_key",
region="gcp-us-central1",
distance_metric="euclidean_squared",
)
assert db._parse_output([_make_row("id1")])[0].score is None
# ── _convert_filters ─────────────────────────────────────────────────