diff --git a/docs/src/js/namespaces/embedding/README.md b/docs/src/js/namespaces/embedding/README.md index 157018e16..a736674c0 100644 --- a/docs/src/js/namespaces/embedding/README.md +++ b/docs/src/js/namespaces/embedding/README.md @@ -25,9 +25,12 @@ ### Type Aliases - [CreateReturnType](type-aliases/CreateReturnType.md) +- [EmbeddingMetadataEntry](type-aliases/EmbeddingMetadataEntry.md) +- [ResolvedEmbeddingFunctionConfig](type-aliases/ResolvedEmbeddingFunctionConfig.md) ### Functions - [LanceSchema](functions/LanceSchema.md) - [getRegistry](functions/getRegistry.md) +- [parseEmbeddingMetadata](functions/parseEmbeddingMetadata.md) - [register](functions/register.md) diff --git a/docs/src/js/namespaces/embedding/functions/parseEmbeddingMetadata.md b/docs/src/js/namespaces/embedding/functions/parseEmbeddingMetadata.md new file mode 100644 index 000000000..d6c381bb3 --- /dev/null +++ b/docs/src/js/namespaces/embedding/functions/parseEmbeddingMetadata.md @@ -0,0 +1,22 @@ +[**@lancedb/lancedb**](../../../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../../../globals.md) / [embedding](../README.md) / parseEmbeddingMetadata + +# Function: parseEmbeddingMetadata() + +```ts +function parseEmbeddingMetadata(json): EmbeddingMetadataEntry[] +``` + +The single parser for `embedding_functions` schema metadata: every reader +goes through here, so the wire contract cannot fork between them. + +## Parameters + +* **json**: `string` + +## Returns + +[`EmbeddingMetadataEntry`](../type-aliases/EmbeddingMetadataEntry.md)[] diff --git a/docs/src/js/namespaces/embedding/type-aliases/EmbeddingMetadataEntry.md b/docs/src/js/namespaces/embedding/type-aliases/EmbeddingMetadataEntry.md new file mode 100644 index 000000000..a1bd247f6 --- /dev/null +++ b/docs/src/js/namespaces/embedding/type-aliases/EmbeddingMetadataEntry.md @@ -0,0 +1,40 @@ +[**@lancedb/lancedb**](../../../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../../../globals.md) / [embedding](../README.md) / EmbeddingMetadataEntry + +# Type Alias: EmbeddingMetadataEntry + +```ts +type EmbeddingMetadataEntry: object; +``` + +One entry of the `embedding_functions` schema metadata, with the column +keys normalized across the bindings' spellings. + +## Type declaration + +### model + +```ts +model: EmbeddingFunction["TOptions"]; +``` + +### name + +```ts +name: string; +``` + +### sourceColumn + +```ts +sourceColumn: string; +``` + +### vectorColumn + +```ts +vectorColumn: string; +``` diff --git a/docs/src/js/namespaces/embedding/type-aliases/ResolvedEmbeddingFunctionConfig.md b/docs/src/js/namespaces/embedding/type-aliases/ResolvedEmbeddingFunctionConfig.md new file mode 100644 index 000000000..864dfedfc --- /dev/null +++ b/docs/src/js/namespaces/embedding/type-aliases/ResolvedEmbeddingFunctionConfig.md @@ -0,0 +1,22 @@ +[**@lancedb/lancedb**](../../../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../../../globals.md) / [embedding](../README.md) / ResolvedEmbeddingFunctionConfig + +# Type Alias: ResolvedEmbeddingFunctionConfig + +```ts +type ResolvedEmbeddingFunctionConfig: EmbeddingFunctionConfig & object; +``` + +An [EmbeddingFunctionConfig] read back from table metadata, where the +vector column is always recorded. + +## Type declaration + +### vectorColumn + +```ts +vectorColumn: string; +``` diff --git a/nodejs/__test__/arrow.test.ts b/nodejs/__test__/arrow.test.ts index 29030d4f8..160f4d5ef 100644 --- a/nodejs/__test__/arrow.test.ts +++ b/nodejs/__test__/arrow.test.ts @@ -173,6 +173,36 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( } describe("The function makeArrowTable", function () { + it("accepts snake_case embedding metadata like camelCase", function () { + const spellings = [ + // biome-ignore lint/style/useNamingConvention: the Python wire spelling + { source_column: "text", vector_column: "vector" }, + { sourceColumn: "text", vectorColumn: "vector" }, + ]; + for (const columns of spellings) { + const schema = new Schema( + [ + new Field("text", new Utf8(), false), + new Field( + "vector", + new FixedSizeList(3, new Field("item", new Float32(), true)), + false, + ), + ], + new Map([ + [ + "embedding_functions", + JSON.stringify([{ name: "mock", model: {}, ...columns }]), + ], + ]), + ); + // The vector field is non-nullable and absent from the data; only a + // recognized embedding config makes that acceptable. + const table = makeArrowTable([{ text: "hello" }], { schema }); + expect(table.numRows).toBe(1); + } + }); + it("will use data types from a provided schema instead of inference", async function () { const schema = new Schema([ new Field("a", new Int32(), false), diff --git a/nodejs/__test__/registry.test.ts b/nodejs/__test__/registry.test.ts index a5cf73e74..973ad8a25 100644 --- a/nodejs/__test__/registry.test.ts +++ b/nodejs/__test__/registry.test.ts @@ -106,6 +106,77 @@ describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => { 'Embedding function with alias "mock-embedding" already exists', ); }); + test("parseFunctions keeps entries sharing a function name", async () => { + class MockEmbeddingFunction extends EmbeddingFunction { + ndims() { + return 3; + } + embeddingDataType() { + return new arrow.Float32() as apiArrow.Float; + } + async computeSourceEmbeddings(data: string[]) { + return data.map(() => [1, 2, 3]); + } + } + register("mock-embedding")(MockEmbeddingFunction); + const parsed = await getRegistry().parseFunctions( + new Map([ + [ + "embedding_functions", + JSON.stringify([ + { + name: "mock-embedding", + sourceColumn: "text", + vectorColumn: "vector_a", + model: {}, + }, + { + name: "mock-embedding", + sourceColumn: "text", + vectorColumn: "vector_b", + model: {}, + }, + ]), + ], + ]), + ); + expect([...parsed.values()].map((f) => f.vectorColumn)).toEqual([ + "vector_a", + "vector_b", + ]); + + // The Python bindings write snake_case keys. + const snake = await getRegistry().parseFunctions( + new Map([ + [ + "embedding_functions", + JSON.stringify([ + { + name: "mock-embedding", + // biome-ignore lint/style/useNamingConvention: the Python wire spelling + source_column: "text", + // biome-ignore lint/style/useNamingConvention: the Python wire spelling + vector_column: "vector_a", + model: {}, + }, + { + name: "mock-embedding", + // biome-ignore lint/style/useNamingConvention: the Python wire spelling + source_column: "text", + // biome-ignore lint/style/useNamingConvention: the Python wire spelling + vector_column: "vector_b", + model: {}, + }, + ]), + ], + ]), + ); + expect([...snake.keys()]).toEqual(["vector_a", "vector_b"]); + expect([...snake.values()].map((f) => f.sourceColumn)).toEqual([ + "text", + "text", + ]); + }); test("schema should contain correct metadata", async () => { class MockEmbeddingFunction extends EmbeddingFunction { constructor(args: FunctionOptions = {}) { diff --git a/nodejs/lancedb/arrow.ts b/nodejs/lancedb/arrow.ts index 587d30b19..8b388d593 100644 --- a/nodejs/lancedb/arrow.ts +++ b/nodejs/lancedb/arrow.ts @@ -48,7 +48,11 @@ import { } from "apache-arrow"; import { Buffers } from "apache-arrow/data"; import { type EmbeddingFunction } from "./embedding/embedding_function"; -import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry"; +import { + EmbeddingFunctionConfig, + getRegistry, + parseEmbeddingMetadata, +} from "./embedding/registry"; import { sanitizeField, sanitizeSchema, @@ -933,7 +937,7 @@ async function applyEmbeddingsFromMetadata( for (const functionEntry of functions.values()) { const sourceColumn = columns[functionEntry.sourceColumn]; - const destColumn = functionEntry.vectorColumn ?? "vector"; + const destColumn = functionEntry.vectorColumn; if (sourceColumn === undefined) { throw new Error( `Cannot apply embedding function because the source column '${functionEntry.sourceColumn}' was not present in the data`, @@ -1385,11 +1389,10 @@ function validateSchemaEmbeddings( // Check schema metadata for embedding functions if (schema.metadata.has("embedding_functions")) { - const embeddings = JSON.parse( + const entries = parseEmbeddingMetadata( schema.metadata.get("embedding_functions")!, ); - // biome-ignore lint/suspicious/noExplicitAny: we don't know the type of `f` - if (embeddings.find((f: any) => f["vectorColumn"] === field.name)) { + if (entries.some((f) => f.vectorColumn === field.name)) { hasEmbeddingFunction = true; } } diff --git a/nodejs/lancedb/embedding/registry.ts b/nodejs/lancedb/embedding/registry.ts index 2eae90ed3..5f32f683c 100644 --- a/nodejs/lancedb/embedding/registry.ts +++ b/nodejs/lancedb/embedding/registry.ts @@ -104,41 +104,29 @@ export class EmbeddingFunctionRegistry { async parseFunctions( this: EmbeddingFunctionRegistry, metadata: Map, - ): Promise> { + ): Promise> { if (!metadata.has("embedding_functions")) { return new Map(); - } else { - type FunctionConfig = { - name: string; - sourceColumn: string; - vectorColumn: string; - model: EmbeddingFunction["TOptions"]; - }; - - const functions = ( - JSON.parse(metadata.get("embedding_functions")!) - ); - - const items: [string, EmbeddingFunctionConfig][] = await Promise.all( - functions.map(async (f) => { - const fn = this.get(f.name); - if (!fn) { - throw new Error(`Function "${f.name}" not found in registry`); - } - const func = await this.get(f.name)!.create(f.model); - return [ - f.name, - { - sourceColumn: f.sourceColumn, - vectorColumn: f.vectorColumn, - function: func, - }, - ]; - }), - ); - - return new Map(items); } + const entries = parseEmbeddingMetadata( + metadata.get("embedding_functions")!, + ); + const items = await Promise.all( + entries.map(async (f): Promise => { + const fn = this.get(f.name); + if (!fn) { + throw new Error(`Function "${f.name}" not found in registry`); + } + const func = await fn.create(f.model); + return { + sourceColumn: f.sourceColumn, + vectorColumn: f.vectorColumn, + function: func, + }; + }), + ); + // Keyed by output column: one function may serve several columns. + return new Map(items.map((config) => [config.vectorColumn, config])); } // biome-ignore lint/suspicious/noExplicitAny: functionToMetadata(conf: EmbeddingFunctionConfig): Record { @@ -218,3 +206,52 @@ export interface EmbeddingFunctionConfig { vectorColumn?: string; function: EmbeddingFunction; } + +/** An [EmbeddingFunctionConfig] read back from table metadata, where the + * vector column is always recorded. */ +export type ResolvedEmbeddingFunctionConfig = EmbeddingFunctionConfig & { + vectorColumn: string; +}; + +/** One entry of the `embedding_functions` schema metadata, with the column + * keys normalized across the bindings' spellings. */ +export type EmbeddingMetadataEntry = { + name: string; + sourceColumn: string; + vectorColumn: string; + model: EmbeddingFunction["TOptions"]; +}; + +/** The single parser for `embedding_functions` schema metadata: every reader + * goes through here, so the wire contract cannot fork between them. */ +export function parseEmbeddingMetadata(json: string): EmbeddingMetadataEntry[] { + // The wire format, honestly: the Python bindings write snake_case keys. + type Raw = { + name: string; + sourceColumn?: string; + // biome-ignore lint/style/useNamingConvention: the Python wire spelling + source_column?: string; + vectorColumn?: string; + // biome-ignore lint/style/useNamingConvention: the Python wire spelling + vector_column?: string; + model: EmbeddingFunction["TOptions"]; + }; + const entries = JSON.parse(json); + const seen = new Set(); + return entries.map((f) => { + const sourceColumn = f.sourceColumn ?? f.source_column; + const vectorColumn = f.vectorColumn ?? f.vector_column; + if (sourceColumn === undefined || vectorColumn === undefined) { + throw new Error( + `Embedding function "${f.name}" metadata names no source or vector column`, + ); + } + if (seen.has(vectorColumn)) { + throw new Error( + `Multiple embedding configs claim vector column "${vectorColumn}"`, + ); + } + seen.add(vectorColumn); + return { name: f.name, sourceColumn, vectorColumn, model: f.model }; + }); +}