fix(ts-oss/turbopuffer): bound euclidean_squared distance to a similarity score (#6580)

This commit is contained in:
Abhay Singh
2026-09-25 17:07:11 +05:30
committed by GitHub
parent 70c676c8c4
commit 5fd01d284a
2 changed files with 92 additions and 1 deletions
@@ -305,7 +305,18 @@ export class TurbopufferDB implements VectorStore {
private parseRows(rows: any[]): VectorStoreResult[] {
return rows.map((row) => {
const { id, $dist, vector, ...rest } = row;
const score = $dist != null ? 1 - $dist : undefined;
let score: number | undefined;
if ($dist == null) {
score = undefined;
} else if (this.distanceMetric === "euclidean_squared") {
// euclidean_squared $dist is an unbounded squared distance, so 1 - $dist
// goes negative for any $dist > 1 and inverts ranking. Convert it to a
// bounded higher-is-better score, mirroring the milvus/baidu stores.
score = 1 / (1 + $dist);
} else {
// Cosine distance is in [0, 2]; 1 - $dist stays a meaningful similarity.
score = 1 - $dist;
}
return { id: String(id), payload: rest, score };
});
}
@@ -0,0 +1,80 @@
/// <reference types="jest" />
/**
* Turbopuffer vector store — score conversion unit tests.
*
* Drives parseRows() through the public search() API with a virtually-mocked
* @turbopuffer/turbopuffer peer, asserting the score returned per metric.
*/
const mockQuery = jest.fn();
jest.mock(
"@turbopuffer/turbopuffer",
() => ({
__esModule: true,
default: class {
namespace() {
return { query: mockQuery };
}
},
}),
{ virtual: true },
);
import { TurbopufferDB } from "../src/vector_stores/turbopuffer";
function makeStore(distanceMetric?: string) {
return new TurbopufferDB({
apiKey: "test-key",
collectionName: "mem0",
...(distanceMetric ? { distanceMetric } : {}),
} as any);
}
async function scoreFor(
distanceMetric: string | undefined,
row: Record<string, any>,
): Promise<number | undefined> {
mockQuery.mockResolvedValueOnce({ rows: [row] });
const results = await makeStore(distanceMetric).search([0.1, 0.2, 0.3], 5);
return results[0].score;
}
describe("TurbopufferDB score conversion", () => {
it("keeps cosine distance as 1 - dist (default metric)", async () => {
// cosine_distance is the default; a distance of 0.25 -> similarity 0.75.
expect(await scoreFor(undefined, { id: "a", $dist: 0.25 })).toBeCloseTo(
0.75,
10,
);
});
it("bounds an unbounded euclidean_squared distance to a 0..1 similarity", async () => {
// $dist = 4.0 is a squared distance. 1 - 4.0 = -3.0 would invert ranking;
// 1 / (1 + 4.0) = 0.2 keeps it higher-is-better and in range.
const score = await scoreFor("euclidean_squared", { id: "a", $dist: 4.0 });
expect(score).toBeCloseTo(0.2, 10);
expect(score!).toBeGreaterThanOrEqual(0);
expect(score!).toBeLessThanOrEqual(1);
});
it("ranks a nearer euclidean_squared hit above a farther one", async () => {
mockQuery.mockResolvedValueOnce({
rows: [
{ id: "near", $dist: 1.0 },
{ id: "far", $dist: 9.0 },
],
});
const results = await makeStore("euclidean_squared").search(
[0.1, 0.2, 0.3],
5,
);
const near = results.find((r) => r.id === "near")!;
const far = results.find((r) => r.id === "far")!;
expect(near.score!).toBeGreaterThan(far.score!);
});
it("preserves an undefined score when the row has no distance", async () => {
expect(await scoreFor("euclidean_squared", { id: "a" })).toBeUndefined();
});
});