@lancedb/lancedb 0.5.0 → 0.5.1

package/lancedb/arrow.ts

@@ -17,10 +17,14 @@ import {
   Binary,
   DataType,
   Field,
+  FixedSizeBinary,
   FixedSizeList,
-  type Float,
+  Float,
   Float32,
+  Int,
+  LargeBinary,
   List,
+  Null,
   RecordBatch,
   RecordBatchFileWriter,
   RecordBatchStreamWriter,
@@ -34,7 +38,99 @@ import {
   vectorFromArray,
 } from "apache-arrow";
 import { type EmbeddingFunction } from "./embedding/embedding_function";
-import { sanitizeSchema } from "./sanitize";
+import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry";
+import { sanitizeField, sanitizeSchema, sanitizeType } from "./sanitize";
+export * from "apache-arrow";
+
+export function isArrowTable(value: object): value is ArrowTable {
+  if (value instanceof ArrowTable) return true;
+  return "schema" in value && "batches" in value;
+}
+
+export function isDataType(value: unknown): value is DataType {
+  return (
+    value instanceof DataType ||
+    DataType.isNull(value) ||
+    DataType.isInt(value) ||
+    DataType.isFloat(value) ||
+    DataType.isBinary(value) ||
+    DataType.isLargeBinary(value) ||
+    DataType.isUtf8(value) ||
+    DataType.isLargeUtf8(value) ||
+    DataType.isBool(value) ||
+    DataType.isDecimal(value) ||
+    DataType.isDate(value) ||
+    DataType.isTime(value) ||
+    DataType.isTimestamp(value) ||
+    DataType.isInterval(value) ||
+    DataType.isDuration(value) ||
+    DataType.isList(value) ||
+    DataType.isStruct(value) ||
+    DataType.isUnion(value) ||
+    DataType.isFixedSizeBinary(value) ||
+    DataType.isFixedSizeList(value) ||
+    DataType.isMap(value) ||
+    DataType.isDictionary(value)
+  );
+}
+export function isNull(value: unknown): value is Null {
+  return value instanceof Null || DataType.isNull(value);
+}
+export function isInt(value: unknown): value is Int {
+  return value instanceof Int || DataType.isInt(value);
+}
+export function isFloat(value: unknown): value is Float {
+  return value instanceof Float || DataType.isFloat(value);
+}
+export function isBinary(value: unknown): value is Binary {
+  return value instanceof Binary || DataType.isBinary(value);
+}
+export function isLargeBinary(value: unknown): value is LargeBinary {
+  return value instanceof LargeBinary || DataType.isLargeBinary(value);
+}
+export function isUtf8(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isUtf8(value);
+}
+export function isLargeUtf8(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isLargeUtf8(value);
+}
+export function isBool(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isBool(value);
+}
+export function isDecimal(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isDecimal(value);
+}
+export function isDate(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isDate(value);
+}
+export function isTime(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isTime(value);
+}
+export function isTimestamp(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isTimestamp(value);
+}
+export function isInterval(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isInterval(value);
+}
+export function isDuration(value: unknown): value is Utf8 {
+  return value instanceof Utf8 || DataType.isDuration(value);
+}
+export function isList(value: unknown): value is List {
+  return value instanceof List || DataType.isList(value);
+}
+export function isStruct(value: unknown): value is Struct {
+  return value instanceof Struct || DataType.isStruct(value);
+}
+export function isUnion(value: unknown): value is Struct {
+  return value instanceof Struct || DataType.isUnion(value);
+}
+export function isFixedSizeBinary(value: unknown): value is FixedSizeBinary {
+  return value instanceof FixedSizeBinary || DataType.isFixedSizeBinary(value);
+}
+
+export function isFixedSizeList(value: unknown): value is FixedSizeList {
+  return value instanceof FixedSizeList || DataType.isFixedSizeList(value);
+}
 
 /** Data type accepted by NodeJS SDK */
 export type Data = Record<string, unknown>[] | ArrowTable;
@@ -198,6 +294,7 @@ export class MakeArrowTableOptions {
 export function makeArrowTable(
   data: Array<Record<string, unknown>>,
   options?: Partial<MakeArrowTableOptions>,
+  metadata?: Map<string, string>,
 ): ArrowTable {
   if (
     data.length === 0 &&
@@ -290,20 +387,41 @@ export function makeArrowTable(
     // `new ArrowTable(schema, batches)` which does not do any schema inference
     const firstTable = new ArrowTable(columns);
     const batchesFixed = firstTable.batches.map(
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
       (batch) => new RecordBatch(opt.schema!, batch.data),
     );
-    return new ArrowTable(opt.schema, batchesFixed);
-  } else {
-    return new ArrowTable(columns);
+    let schema: Schema;
+    if (metadata !== undefined) {
+      let schemaMetadata = opt.schema.metadata;
+      if (schemaMetadata.size === 0) {
+        schemaMetadata = metadata;
+      } else {
+        for (const [key, entry] of schemaMetadata.entries()) {
+          schemaMetadata.set(key, entry);
+        }
+      }
+
+      schema = new Schema(opt.schema.fields, schemaMetadata);
+    } else {
+      schema = opt.schema;
+    }
+    return new ArrowTable(schema, batchesFixed);
   }
+  const tbl = new ArrowTable(columns);
+  if (metadata !== undefined) {
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+    (<any>tbl.schema).metadata = metadata;
+  }
+  return tbl;
 }
 
 /**
  * Create an empty Arrow table with the provided schema
  */
-export function makeEmptyTable(schema: Schema): ArrowTable {
-  return makeArrowTable([], { schema });
+export function makeEmptyTable(
+  schema: Schema,
+  metadata?: Map<string, string>,
+): ArrowTable {
+  return makeArrowTable([], { schema }, metadata);
 }
 
 /**
@@ -375,13 +493,75 @@ function makeVector(
   }
 }
 
+/** Helper function to apply embeddings from metadata to an input table */
+async function applyEmbeddingsFromMetadata(
+  table: ArrowTable,
+  schema: Schema,
+): Promise<ArrowTable> {
+  const registry = getRegistry();
+  const functions = registry.parseFunctions(schema.metadata);
+
+  const columns = Object.fromEntries(
+    table.schema.fields.map((field) => [
+      field.name,
+      table.getChild(field.name)!,
+    ]),
+  );
+
+  for (const functionEntry of functions.values()) {
+    const sourceColumn = columns[functionEntry.sourceColumn];
+    const destColumn = functionEntry.vectorColumn ?? "vector";
+    if (sourceColumn === undefined) {
+      throw new Error(
+        `Cannot apply embedding function because the source column '${functionEntry.sourceColumn}' was not present in the data`,
+      );
+    }
+    if (columns[destColumn] !== undefined) {
+      throw new Error(
+        `Attempt to apply embeddings to table failed because column ${destColumn} already existed`,
+      );
+    }
+    if (table.batches.length > 1) {
+      throw new Error(
+        "Internal error: `makeArrowTable` unexpectedly created a table with more than one batch",
+      );
+    }
+    const values = sourceColumn.toArray();
+
+    const vectors =
+      await functionEntry.function.computeSourceEmbeddings(values);
+    if (vectors.length !== values.length) {
+      throw new Error(
+        "Embedding function did not return an embedding for each input element",
+      );
+    }
+    let destType: DataType;
+    const dtype = schema.fields.find((f) => f.name === destColumn)!.type;
+    if (isFixedSizeList(dtype)) {
+      destType = sanitizeType(dtype);
+    } else {
+      throw new Error(
+        "Expected FixedSizeList as datatype for vector field, instead got: " +
+          dtype,
+      );
+    }
+
+    const vector = makeVector(vectors, destType);
+    columns[destColumn] = vector;
+  }
+  const newTable = new ArrowTable(columns);
+  return alignTable(newTable, schema);
+}
+
 /** Helper function to apply embeddings to an input table */
 async function applyEmbeddings<T>(
   table: ArrowTable,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<ArrowTable> {
-  if (embeddings == null) {
+  if (schema?.metadata.has("embedding_functions")) {
+    return applyEmbeddingsFromMetadata(table, schema!);
+  } else if (embeddings == null || embeddings === undefined) {
     return table;
   }
 
@@ -399,8 +579,9 @@ async function applyEmbeddings<T>(
   const newColumns = Object.fromEntries(colEntries);
 
   const sourceColumn = newColumns[embeddings.sourceColumn];
-  const destColumn = embeddings.destColumn ?? "vector";
-  const innerDestType = embeddings.embeddingDataType ?? new Float32();
+  const destColumn = embeddings.vectorColumn ?? "vector";
+  const innerDestType =
+    embeddings.function.embeddingDataType() ?? new Float32();
   if (sourceColumn === undefined) {
     throw new Error(
       `Cannot apply embedding function because the source column '${embeddings.sourceColumn}' was not present in the data`,
@@ -414,11 +595,9 @@ async function applyEmbeddings<T>(
       // if we call convertToTable with 0 records and a schema that includes the embedding
       return table;
     }
-    if (embeddings.embeddingDimension !== undefined) {
-      const destType = newVectorType(
-        embeddings.embeddingDimension,
-        innerDestType,
-      );
+    const dimensions = embeddings.function.ndims();
+    if (dimensions !== undefined) {
+      const destType = newVectorType(dimensions, innerDestType);
       newColumns[destColumn] = makeVector([], destType);
     } else if (schema != null) {
       const destField = schema.fields.find((f) => f.name === destColumn);
@@ -446,7 +625,9 @@ async function applyEmbeddings<T>(
       );
     }
     const values = sourceColumn.toArray();
-    const vectors = await embeddings.embed(values as T[]);
+    const vectors = await embeddings.function.computeSourceEmbeddings(
+      values as T[],
+    );
     if (vectors.length !== values.length) {
       throw new Error(
         "Embedding function did not return an embedding for each input element",
@@ -486,9 +667,9 @@ async function applyEmbeddings<T>(
  * embedding columns.  If no schema is provded then embedding columns will
  * be placed at the end of the table, after all of the input columns.
  */
-export async function convertToTable<T>(
+export async function convertToTable(
   data: Array<Record<string, unknown>>,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   makeTableOptions?: Partial<MakeArrowTableOptions>,
 ): Promise<ArrowTable> {
   const table = makeArrowTable(data, makeTableOptions);
@@ -496,13 +677,13 @@ export async function convertToTable<T>(
 }
 
 /** Creates the Arrow Type for a Vector column with dimension `dim` */
-function newVectorType<T extends Float>(
+export function newVectorType<T extends Float>(
   dim: number,
   innerType: T,
 ): FixedSizeList<T> {
   // in Lance we always default to have the elements nullable, so we need to set it to true
   // otherwise we often get schema mismatches because the stored data always has schema with nullable elements
-  const children = new Field<T>("item", innerType, true);
+  const children = new Field("item", <T>sanitizeType(innerType), true);
   return new FixedSizeList(dim, children);
 }
 
@@ -513,9 +694,9 @@ function newVectorType<T extends Float>(
  *
  * `schema` is required if data is empty
  */
-export async function fromRecordsToBuffer<T>(
+export async function fromRecordsToBuffer(
   data: Array<Record<string, unknown>>,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<Buffer> {
   if (schema !== undefined && schema !== null) {
@@ -533,9 +714,9 @@ export async function fromRecordsToBuffer<T>(
  *
  * `schema` is required if data is empty
  */
-export async function fromRecordsToStreamBuffer<T>(
+export async function fromRecordsToStreamBuffer(
   data: Array<Record<string, unknown>>,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<Buffer> {
   if (schema !== undefined && schema !== null) {
@@ -554,9 +735,9 @@ export async function fromRecordsToStreamBuffer<T>(
  *
  * `schema` is required if the table is empty
  */
-export async function fromTableToBuffer<T>(
+export async function fromTableToBuffer(
   table: ArrowTable,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<Buffer> {
   if (schema !== undefined && schema !== null) {
@@ -575,19 +756,19 @@ export async function fromTableToBuffer<T>(
  *
  * `schema` is required if the table is empty
  */
-export async function fromDataToBuffer<T>(
+export async function fromDataToBuffer(
   data: Data,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<Buffer> {
   if (schema !== undefined && schema !== null) {
     schema = sanitizeSchema(schema);
   }
-  if (data instanceof ArrowTable) {
+  if (isArrowTable(data)) {
     return fromTableToBuffer(data, embeddings, schema);
   } else {
-    const table = await convertToTable(data);
-    return fromTableToBuffer(table, embeddings, schema);
+    const table = await convertToTable(data, embeddings, { schema });
+    return fromTableToBuffer(table);
   }
 }
 
@@ -599,9 +780,9 @@ export async function fromDataToBuffer<T>(
  *
  * `schema` is required if the table is empty
  */
-export async function fromTableToStreamBuffer<T>(
+export async function fromTableToStreamBuffer(
   table: ArrowTable,
-  embeddings?: EmbeddingFunction<T>,
+  embeddings?: EmbeddingFunctionConfig,
   schema?: Schema,
 ): Promise<Buffer> {
   const tableWithEmbeddings = await applyEmbeddings(table, embeddings, schema);
@@ -664,10 +845,25 @@ function validateSchemaEmbeddings(
   // if it does not, we add it to the list of missing embedding fields
   // Finally, we check if those missing embedding fields are `this._embeddings`
   // if they are not, we throw an error
-  for (const field of schema.fields) {
-    if (field.type instanceof FixedSizeList) {
+  for (let field of schema.fields) {
+    if (isFixedSizeList(field.type)) {
+      field = sanitizeField(field);
+
       if (data.length !== 0 && data?.[0]?.[field.name] === undefined) {
-        missingEmbeddingFields.push(field);
+        if (schema.metadata.has("embedding_functions")) {
+          const embeddings = JSON.parse(
+            schema.metadata.get("embedding_functions")!,
+          );
+          if (
+            // biome-ignore lint/suspicious/noExplicitAny: we don't know the type of `f`
+            embeddings.find((f: any) => f["vectorColumn"] === field.name) ===
+            undefined
+          ) {
+            missingEmbeddingFields.push(field);
+          }
+        } else {
+          missingEmbeddingFields.push(field);
+        }
       } else {
         fields.push(field);
       }

package/lancedb/connection.ts

@@ -12,8 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { Table as ArrowTable, Schema } from "apache-arrow";
-import { fromTableToBuffer, makeArrowTable, makeEmptyTable } from "./arrow";
+import { Table as ArrowTable, Schema } from "./arrow";
+import {
+  fromTableToBuffer,
+  isArrowTable,
+  makeArrowTable,
+  makeEmptyTable,
+} from "./arrow";
+import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry";
 import { ConnectionOptions, Connection as LanceDbConnection } from "./native";
 import { Table } from "./table";
 
@@ -65,6 +71,8 @@ export interface CreateTableOptions {
    * The available options are described at https://lancedb.github.io/lancedb/guides/storage/
    */
   storageOptions?: Record<string, string>;
+  schema?: Schema;
+  embeddingFunction?: EmbeddingFunctionConfig;
 }
 
 export interface OpenTableOptions {
@@ -174,6 +182,7 @@ export class Connection {
       cleanseStorageOptions(options?.storageOptions),
       options?.indexCacheSize,
     );
+
     return new Table(innerTable);
   }
 
@@ -196,18 +205,24 @@ export class Connection {
     }
 
     let table: ArrowTable;
-    if (data instanceof ArrowTable) {
+    if (isArrowTable(data)) {
       table = data;
     } else {
-      table = makeArrowTable(data);
+      table = makeArrowTable(data, options);
     }
-    const buf = await fromTableToBuffer(table);
+
+    const buf = await fromTableToBuffer(
+      table,
+      options?.embeddingFunction,
+      options?.schema,
+    );
     const innerTable = await this.inner.createTable(
       name,
       buf,
       mode,
       cleanseStorageOptions(options?.storageOptions),
     );
+
     return new Table(innerTable);
   }
 
@@ -227,8 +242,14 @@ export class Connection {
     if (mode === "create" && existOk) {
       mode = "exist_ok";
     }
+    let metadata: Map<string, string> | undefined = undefined;
+    if (options?.embeddingFunction !== undefined) {
+      const embeddingFunction = options.embeddingFunction;
+      const registry = getRegistry();
+      metadata = registry.getTableMetadata([embeddingFunction]);
+    }
 
-    const table = makeEmptyTable(schema);
+    const table = makeEmptyTable(schema, metadata);
     const buf = await fromTableToBuffer(table);
     const innerTable = await this.inner.createEmptyTable(
       name,

package/lancedb/embedding/embedding_function.ts

@@ -1,4 +1,4 @@
-// Copyright 2023 Lance Developers.
+// Copyright 2024 Lance Developers.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -12,67 +12,151 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { type Float } from "apache-arrow";
+import "reflect-metadata";
+import {
+  DataType,
+  Field,
+  FixedSizeList,
+  Float,
+  Float32,
+  isDataType,
+  isFixedSizeList,
+  isFloat,
+  newVectorType,
+} from "../arrow";
+import { sanitizeType } from "../sanitize";
 
 /**
- * An embedding function that automatically creates vector representation for a given column.
+ * Options for a given embedding function
  */
-export interface EmbeddingFunction<T> {
-  /**
-   * The name of the column that will be used as input for the Embedding Function.
-   */
-  sourceColumn: string;
+export interface FunctionOptions {
+  // biome-ignore lint/suspicious/noExplicitAny: options can be anything
+  [key: string]: any;
+}
 
+/**
+ * An embedding function that automatically creates vector representation for a given column.
+ */
+export abstract class EmbeddingFunction<
+  // biome-ignore lint/suspicious/noExplicitAny: we don't know what the implementor will do
+  T = any,
+  M extends FunctionOptions = FunctionOptions,
+> {
   /**
-   * The data type of the embedding
+   * Convert the embedding function to a JSON object
+   * It is used to serialize the embedding function to the schema
+   * It's important that any object returned by this method contains all the necessary
+   * information to recreate the embedding function
    *
-   * The embedding function should return `number`.  This will be converted into
-   * an Arrow float array.  By default this will be Float32 but this property can
-   * be used to control the conversion.
-   */
-  embeddingDataType?: Float;
-
-  /**
-   * The dimension of the embedding
+   * It should return the same object that was passed to the constructor
+   * If it does not, the embedding function will not be able to be recreated, or could be recreated incorrectly
    *
-   * This is optional, normally this can be determined by looking at the results of
-   * `embed`.  If this is not specified, and there is an attempt to apply the embedding
-   * to an empty table, then that process will fail.
+   * @example
+   * ```ts
+   * class MyEmbeddingFunction extends EmbeddingFunction {
+   *   constructor(options: {model: string, timeout: number}) {
+   *     super();
+   *     this.model = options.model;
+   *     this.timeout = options.timeout;
+   *   }
+   *   toJSON() {
+   *     return {
+   *       model: this.model,
+   *       timeout: this.timeout,
+   *     };
+   * }
+   * ```
    */
-  embeddingDimension?: number;
+  abstract toJSON(): Partial<M>;
 
   /**
-   * The name of the column that will contain the embedding
+   * sourceField is used in combination with `LanceSchema` to provide a declarative data model
    *
-   * By default this is "vector"
+   * @param optionsOrDatatype - The options for the field or the datatype
+   *
+   * @see {@link lancedb.LanceSchema}
    */
-  destColumn?: string;
+  sourceField(
+    optionsOrDatatype: Partial<FieldOptions> | DataType,
+  ): [DataType, Map<string, EmbeddingFunction>] {
+    let datatype = isDataType(optionsOrDatatype)
+      ? optionsOrDatatype
+      : optionsOrDatatype?.datatype;
+    if (!datatype) {
+      throw new Error("Datatype is required");
+    }
+    datatype = sanitizeType(datatype);
+    const metadata = new Map<string, EmbeddingFunction>();
+    metadata.set("source_column_for", this);
+
+    return [datatype, metadata];
+  }
 
   /**
-   * Should the source column be excluded from the resulting table
+   * vectorField is used in combination with `LanceSchema` to provide a declarative data model
    *
-   * By default the source column is included.  Set this to true and
-   * only the embedding will be stored.
+   * @param options - The options for the field
+   *
+   * @see {@link lancedb.LanceSchema}
    */
-  excludeSource?: boolean;
+  vectorField(
+    options?: Partial<FieldOptions>,
+  ): [DataType, Map<string, EmbeddingFunction>] {
+    let dtype: DataType;
+    const dims = this.ndims() ?? options?.dims;
+    if (!options?.datatype) {
+      if (dims === undefined) {
+        throw new Error("ndims is required for vector field");
+      }
+      dtype = new FixedSizeList(dims, new Field("item", new Float32(), true));
+    } else {
+      if (isFixedSizeList(options.datatype)) {
+        dtype = options.datatype;
+      } else if (isFloat(options.datatype)) {
+        if (dims === undefined) {
+          throw new Error("ndims is required for vector field");
+        }
+        dtype = newVectorType(dims, options.datatype);
+      } else {
+        throw new Error(
+          "Expected FixedSizeList or Float as datatype for vector field",
+        );
+      }
+    }
+    const metadata = new Map<string, EmbeddingFunction>();
+    metadata.set("vector_column_for", this);
+
+    return [dtype, metadata];
+  }
+
+  /** The number of dimensions of the embeddings */
+  ndims(): number | undefined {
+    return undefined;
+  }
+
+  /** The datatype of the embeddings */
+  abstract embeddingDataType(): Float;
 
   /**
    * Creates a vector representation for the given values.
    */
-  embed: (data: T[]) => Promise<number[][]>;
-}
+  abstract computeSourceEmbeddings(
+    data: T[],
+  ): Promise<number[][] | Float32Array[] | Float64Array[]>;
 
-/** Test if the input seems to be an embedding function */
-export function isEmbeddingFunction<T>(
-  value: unknown,
-): value is EmbeddingFunction<T> {
-  if (typeof value !== "object" || value === null) {
-    return false;
-  }
-  if (!("sourceColumn" in value) || !("embed" in value)) {
-    return false;
+  /**
+  Compute the embeddings for a single query
+ */
+  async computeQueryEmbeddings(
+    data: T,
+  ): Promise<number[] | Float32Array | Float64Array> {
+    return this.computeSourceEmbeddings([data]).then(
+      (embeddings) => embeddings[0],
+    );
   }
-  return (
-    typeof value.sourceColumn === "string" && typeof value.embed === "function"
-  );
+}
+
+export interface FieldOptions<T extends DataType = DataType> {
+  datatype: T;
+  dims?: number;
 }