perf(ts-oss): cut Oracle vector store round trips per review feedback (#6835)

This commit is contained in:
Kartik
2026-08-06 15:12:14 +05:30
committed by GitHub
parent 3f39fba28f
commit beea626f0a
2 changed files with 108 additions and 33 deletions
+26 -18
View File
@@ -1,4 +1,4 @@
import type { Connection, Pool } from "oracledb";
import type { BindParameters, Connection, Pool } from "oracledb";
import { v4 as uuidv4 } from "uuid";
import { VectorStore } from "./base";
import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types";
@@ -567,17 +567,24 @@ export class OracleAIVectorSearch implements VectorStore {
): Promise<void> {
await this.initialize();
if (vectors.length === 0) return;
await this.withConnection(async (connection) => {
for (let i = 0; i < vectors.length; i++) {
await connection.execute(
`INSERT INTO ${this.collectionName} (id, vector, payload) VALUES (:id, :vector, :payload)`,
{
id: ids[i],
vector: this.vectorBind(vectors[i]),
payload: this.payloadBind(payloads[i] ?? {}),
await connection.executeMany(
`INSERT INTO ${this.collectionName} (id, vector, payload) VALUES (:id, :vector, :payload)`,
vectors.map((vector, i) => ({
id: ids[i],
vector: new Float32Array(vector),
payload: payloads[i] ?? {},
})) as BindParameters[],
{
bindDefs: {
id: { type: this.oracledb.DB_TYPE_VARCHAR, maxSize: 36 },
vector: { type: this.oracledb.DB_TYPE_VECTOR },
payload: { type: this.oracledb.DB_TYPE_JSON },
},
);
}
},
);
}, true);
}
@@ -589,8 +596,12 @@ export class OracleAIVectorSearch implements VectorStore {
await this.initialize();
const [whereClause, filterBinds] = buildWhereClause(filters);
const hasFilter = whereClause.length > 0;
const selectClause = hasFilter
? "SELECT"
: `SELECT /*+ VECTOR_INDEX_TRANSFORM(${this.collectionName}) */`;
const sql =
`SELECT id, payload, VECTOR_DISTANCE(vector, :query_vec, ${this.distanceMetric}) distance ` +
`${selectClause} id, payload, VECTOR_DISTANCE(vector, :query_vec, ${this.distanceMetric}) distance ` +
`FROM ${this.collectionName} ${whereClause} ORDER BY distance FETCH APPROX FIRST :max_rows ROWS ONLY`;
const rows = await this.withConnection(async (connection) => {
@@ -687,20 +698,17 @@ export class OracleAIVectorSearch implements VectorStore {
return this.withConnection(async (connection) => {
const listResult = await connection.execute<any[]>(
`SELECT id, payload FROM ${this.collectionName} ${whereClause} FETCH FIRST :max_rows ROWS ONLY`,
`SELECT id, payload, COUNT(*) OVER () total FROM ${this.collectionName} ${whereClause} FETCH FIRST :max_rows ROWS ONLY`,
{ ...filterBinds, max_rows: topK },
);
const countResult = await connection.execute<any[]>(
`SELECT COUNT(*) FROM ${this.collectionName} ${whereClause}`,
filterBinds,
);
const results = (listResult.rows ?? []).map((row) => ({
const rows = listResult.rows ?? [];
const results = rows.map((row) => ({
id: row[0],
payload: this.loadPayload(row[1]),
}));
return [results, Number(countResult.rows?.[0]?.[0] ?? 0)];
return [results, Number(rows[0]?.[2] ?? 0)];
});
}
+82 -15
View File
@@ -2,10 +2,13 @@
/** Oracle AI Vector Search filter, config and SQL tests. The driver is mocked, so no database is needed. */
const DB_TYPE_VECTOR = { name: "DB_TYPE_VECTOR" };
const DB_TYPE_JSON = { name: "DB_TYPE_JSON" };
const DB_TYPE_VARCHAR = { name: "DB_TYPE_VARCHAR" };
jest.mock("oracledb", () => ({ thin: true, DB_TYPE_VECTOR, DB_TYPE_JSON }), {
virtual: true,
});
jest.mock(
"oracledb",
() => ({ thin: true, DB_TYPE_VECTOR, DB_TYPE_JSON, DB_TYPE_VARCHAR }),
{ virtual: true },
);
import {
OracleAIVectorSearch,
@@ -13,7 +16,7 @@ import {
quoteIdentifier,
} from "../src/vector_stores/oracledb";
type Call = { sql: string; binds: any };
type Call = { sql: string; binds: any; options?: any };
function fakeConnection(calls: Call[], resultsBySql: Array<any[][]>) {
let selectIndex = 0;
@@ -26,6 +29,10 @@ function fakeConnection(calls: Call[], resultsBySql: Array<any[][]>) {
}
return { rows: [] };
},
async executeMany(sql: string, binds: any[], options: any = {}) {
calls.push({ sql: sql.replace(/\s+/g, " ").trim(), binds, options });
return { rows: [] };
},
async commit() {},
async rollback() {},
async close() {},
@@ -249,13 +256,45 @@ describe("OracleAIVectorSearch SQL", () => {
await makeStore(calls).insert([[1, 2, 3]], ["id-1"], [{ data: "hello" }]);
const insert = calls.find((c) => c.sql.startsWith("INSERT INTO"))!;
expect(insert.binds.id).toBe("id-1");
expect(insert.binds.vector.type).toBe(DB_TYPE_VECTOR);
expect(insert.binds.vector.val).toEqual(new Float32Array([1, 2, 3]));
expect(insert.binds.payload).toEqual({
type: DB_TYPE_JSON,
val: { data: "hello" },
expect(insert.binds).toEqual([
{
id: "id-1",
vector: new Float32Array([1, 2, 3]),
payload: { data: "hello" },
},
]);
expect(insert.options.bindDefs.id).toEqual({
type: DB_TYPE_VARCHAR,
maxSize: 36,
});
expect(insert.options.bindDefs.vector.type).toBe(DB_TYPE_VECTOR);
expect(insert.options.bindDefs.payload.type).toBe(DB_TYPE_JSON);
});
it("issues a single executeMany call with one bind row per vector on a multi-row insert", async () => {
const calls: Call[] = [];
await makeStore(calls).insert(
[
[1, 2, 3],
[4, 5, 6],
],
["id-1", "id-2"],
[{ a: 1 }, { b: 2 }],
);
const inserts = calls.filter((c) => c.sql.startsWith("INSERT INTO"));
expect(inserts).toHaveLength(1);
expect(inserts[0].binds).toHaveLength(2);
expect(inserts[0].binds[0].id).toBe("id-1");
expect(inserts[0].binds[1].id).toBe("id-2");
});
it("issues no insert statement when inserting an empty batch", async () => {
const calls: Call[] = [];
await makeStore(calls).insert([], [], []);
expect(calls.filter((c) => c.sql.startsWith("INSERT INTO"))).toHaveLength(
0,
);
});
it("converts cosine distance to a similarity score", async () => {
@@ -266,7 +305,7 @@ describe("OracleAIVectorSearch SQL", () => {
expect(results).toEqual([
{ id: "id-1", payload: { data: "hello" }, score: 0.75 },
]);
const select = calls.find((c) => c.sql.startsWith("SELECT id, payload,"))!;
const select = calls.find((c) => c.sql.includes("VECTOR_DISTANCE"))!;
expect(select.sql).toContain(
"VECTOR_DISTANCE(vector, :query_vec, COSINE) distance",
);
@@ -282,6 +321,24 @@ describe("OracleAIVectorSearch SQL", () => {
expect(result.score).toBeCloseTo(0.4);
});
it("adds the VECTOR_INDEX_TRANSFORM hint when searching without filters", async () => {
const calls: Call[] = [];
const store = makeStore(calls, [[["id-1", {}, 0.1]]]);
await store.search([1, 2, 3], 5);
const select = calls.find((c) => c.sql.includes("VECTOR_DISTANCE"))!;
expect(select.sql).toContain('/*+ VECTOR_INDEX_TRANSFORM("mem0") */');
});
it("omits the VECTOR_INDEX_TRANSFORM hint when searching with filters", async () => {
const calls: Call[] = [];
const store = makeStore(calls, [[["id-1", {}, 0.1]]]);
await store.search([1, 2, 3], 5, { user_id: "alice" });
const select = calls.find((c) => c.sql.includes("VECTOR_DISTANCE"))!;
expect(select.sql).not.toContain("VECTOR_INDEX_TRANSFORM");
});
it("parses a payload returned as a JSON string", async () => {
const store = makeStore([], [[["id-1", '{"data":"hello"}']]]);
expect(await store.get("id-1")).toEqual({
@@ -313,16 +370,26 @@ describe("OracleAIVectorSearch SQL", () => {
it("returns rows and the total count from list", async () => {
const calls: Call[] = [];
const store = makeStore(calls, [[["id-1", { data: "hello" }]], [[7]]]);
const store = makeStore(calls, [[["id-1", { data: "hello" }, 7]]]);
const [results, count] = await store.list({ user_id: "alice" }, 10);
expect(results).toEqual([{ id: "id-1", payload: { data: "hello" } }]);
expect(count).toBe(7);
const list = calls.find((c) =>
c.sql.startsWith("SELECT id, payload FROM"),
)!;
const selects = calls.filter((c) => /^SELECT/i.test(c.sql));
expect(selects).toHaveLength(1);
const [list] = selects;
expect(list.sql).toContain(
"SELECT id, payload, COUNT(*) OVER () total FROM",
);
expect(list.sql).toContain("WHERE JSON_EXISTS(payload,");
expect(list.binds).toEqual({ f_0: "alice", max_rows: 10 });
});
it("returns a total of 0 from list when no rows match", async () => {
const store = makeStore([], [[]]);
const [results, count] = await store.list();
expect(results).toEqual([]);
expect(count).toBe(0);
});
});