@lancedb/lancedb 0.5.0 → 0.5.2

package/nodejs-artifacts/arrow.js

@@ -12,10 +12,133 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __exportStar = (this && this.__exportStar) || function(m, exports) {
+    for (var p in m) if (p !== "default" && !Object.prototype.hasOwnProperty.call(exports, p)) __createBinding(exports, m, p);
+};
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.createEmptyTable = exports.fromTableToStreamBuffer = exports.fromDataToBuffer = exports.fromTableToBuffer = exports.fromRecordsToStreamBuffer = exports.fromRecordsToBuffer = exports.convertToTable = exports.makeEmptyTable = exports.makeArrowTable = exports.MakeArrowTableOptions = exports.VectorColumnOptions = void 0;
+exports.createEmptyTable = exports.fromTableToStreamBuffer = exports.fromDataToBuffer = exports.fromTableToBuffer = exports.fromRecordsToStreamBuffer = exports.fromRecordsToBuffer = exports.newVectorType = exports.convertToTable = exports.makeEmptyTable = exports.makeArrowTable = exports.MakeArrowTableOptions = exports.VectorColumnOptions = exports.isFixedSizeList = exports.isFixedSizeBinary = exports.isUnion = exports.isStruct = exports.isList = exports.isDuration = exports.isInterval = exports.isTimestamp = exports.isTime = exports.isDate = exports.isDecimal = exports.isBool = exports.isLargeUtf8 = exports.isUtf8 = exports.isLargeBinary = exports.isBinary = exports.isFloat = exports.isInt = exports.isNull = exports.isDataType = exports.isArrowTable = void 0;
 const apache_arrow_1 = require("apache-arrow");
+const registry_1 = require("./embedding/registry");
 const sanitize_1 = require("./sanitize");
+__exportStar(require("apache-arrow"), exports);
+function isArrowTable(value) {
+    if (value instanceof apache_arrow_1.Table)
+        return true;
+    return "schema" in value && "batches" in value;
+}
+exports.isArrowTable = isArrowTable;
+function isDataType(value) {
+    return (value instanceof apache_arrow_1.DataType ||
+        apache_arrow_1.DataType.isNull(value) ||
+        apache_arrow_1.DataType.isInt(value) ||
+        apache_arrow_1.DataType.isFloat(value) ||
+        apache_arrow_1.DataType.isBinary(value) ||
+        apache_arrow_1.DataType.isLargeBinary(value) ||
+        apache_arrow_1.DataType.isUtf8(value) ||
+        apache_arrow_1.DataType.isLargeUtf8(value) ||
+        apache_arrow_1.DataType.isBool(value) ||
+        apache_arrow_1.DataType.isDecimal(value) ||
+        apache_arrow_1.DataType.isDate(value) ||
+        apache_arrow_1.DataType.isTime(value) ||
+        apache_arrow_1.DataType.isTimestamp(value) ||
+        apache_arrow_1.DataType.isInterval(value) ||
+        apache_arrow_1.DataType.isDuration(value) ||
+        apache_arrow_1.DataType.isList(value) ||
+        apache_arrow_1.DataType.isStruct(value) ||
+        apache_arrow_1.DataType.isUnion(value) ||
+        apache_arrow_1.DataType.isFixedSizeBinary(value) ||
+        apache_arrow_1.DataType.isFixedSizeList(value) ||
+        apache_arrow_1.DataType.isMap(value) ||
+        apache_arrow_1.DataType.isDictionary(value));
+}
+exports.isDataType = isDataType;
+function isNull(value) {
+    return value instanceof apache_arrow_1.Null || apache_arrow_1.DataType.isNull(value);
+}
+exports.isNull = isNull;
+function isInt(value) {
+    return value instanceof apache_arrow_1.Int || apache_arrow_1.DataType.isInt(value);
+}
+exports.isInt = isInt;
+function isFloat(value) {
+    return value instanceof apache_arrow_1.Float || apache_arrow_1.DataType.isFloat(value);
+}
+exports.isFloat = isFloat;
+function isBinary(value) {
+    return value instanceof apache_arrow_1.Binary || apache_arrow_1.DataType.isBinary(value);
+}
+exports.isBinary = isBinary;
+function isLargeBinary(value) {
+    return value instanceof apache_arrow_1.LargeBinary || apache_arrow_1.DataType.isLargeBinary(value);
+}
+exports.isLargeBinary = isLargeBinary;
+function isUtf8(value) {
+    return value instanceof apache_arrow_1.Utf8 || apache_arrow_1.DataType.isUtf8(value);
+}
+exports.isUtf8 = isUtf8;
+function isLargeUtf8(value) {
+    return value instanceof apache_arrow_1.Utf8 || apache_arrow_1.DataType.isLargeUtf8(value);
+}
+exports.isLargeUtf8 = isLargeUtf8;
+function isBool(value) {
+    return value instanceof apache_arrow_1.Utf8 || apache_arrow_1.DataType.isBool(value);
+}
+exports.isBool = isBool;
+function isDecimal(value) {
+    return value instanceof apache_arrow_1.Utf8 || apache_arrow_1.DataType.isDecimal(value);
+}
+exports.isDecimal = isDecimal;
+function isDate(value) {
+    return value instanceof apache_arrow_1.Utf8 || apache_arrow_1.DataType.isDate(value);
+}
+exports.isDate = isDate;
+function isTime(value) {
+    return value instanceof apache_arrow_1.Utf8 || apache_arrow_1.DataType.isTime(value);
+}
+exports.isTime = isTime;
+function isTimestamp(value) {
+    return value instanceof apache_arrow_1.Utf8 || apache_arrow_1.DataType.isTimestamp(value);
+}
+exports.isTimestamp = isTimestamp;
+function isInterval(value) {
+    return value instanceof apache_arrow_1.Utf8 || apache_arrow_1.DataType.isInterval(value);
+}
+exports.isInterval = isInterval;
+function isDuration(value) {
+    return value instanceof apache_arrow_1.Utf8 || apache_arrow_1.DataType.isDuration(value);
+}
+exports.isDuration = isDuration;
+function isList(value) {
+    return value instanceof apache_arrow_1.List || apache_arrow_1.DataType.isList(value);
+}
+exports.isList = isList;
+function isStruct(value) {
+    return value instanceof apache_arrow_1.Struct || apache_arrow_1.DataType.isStruct(value);
+}
+exports.isStruct = isStruct;
+function isUnion(value) {
+    return value instanceof apache_arrow_1.Struct || apache_arrow_1.DataType.isUnion(value);
+}
+exports.isUnion = isUnion;
+function isFixedSizeBinary(value) {
+    return value instanceof apache_arrow_1.FixedSizeBinary || apache_arrow_1.DataType.isFixedSizeBinary(value);
+}
+exports.isFixedSizeBinary = isFixedSizeBinary;
+function isFixedSizeList(value) {
+    return value instanceof apache_arrow_1.FixedSizeList || apache_arrow_1.DataType.isFixedSizeList(value);
+}
+exports.isFixedSizeList = isFixedSizeList;
 /*
  * Options to control how a column should be converted to a vector array
  */
@@ -61,6 +184,7 @@ class MakeArrowTableOptions {
         vector: new VectorColumnOptions(),
     };
     embeddings;
+    embeddingFunction;
     /**
      * If true then string columns will be encoded with dictionary encoding
      *
@@ -168,7 +292,7 @@ exports.MakeArrowTableOptions = MakeArrowTableOptions;
  * assert.deepEqual(table.schema, schema)
  * ```
  */
-function makeArrowTable(data, options) {
+function makeArrowTable(data, options, metadata) {
     if (data.length === 0 &&
         (options?.schema === undefined || options?.schema === null)) {
         throw new Error("At least one record or a schema needs to be provided");
@@ -176,7 +300,7 @@ function makeArrowTable(data, options) {
     const opt = new MakeArrowTableOptions(options !== undefined ? options : {});
     if (opt.schema !== undefined && opt.schema !== null) {
         opt.schema = (0, sanitize_1.sanitizeSchema)(opt.schema);
-        opt.schema = validateSchemaEmbeddings(opt.schema, data, opt.embeddings);
+        opt.schema = validateSchemaEmbeddings(opt.schema, data, options?.embeddingFunction);
     }
     const columns = {};
     // TODO: sample dataset to find missing columns
@@ -246,21 +370,38 @@ function makeArrowTable(data, options) {
         // then patch the schema of the batches so we can use
         // `new ArrowTable(schema, batches)` which does not do any schema inference
         const firstTable = new apache_arrow_1.Table(columns);
-        const batchesFixed = firstTable.batches.map(
-        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-        (batch) => new apache_arrow_1.RecordBatch(opt.schema, batch.data));
-        return new apache_arrow_1.Table(opt.schema, batchesFixed);
+        const batchesFixed = firstTable.batches.map((batch) => new apache_arrow_1.RecordBatch(opt.schema, batch.data));
+        let 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 apache_arrow_1.Schema(opt.schema.fields, schemaMetadata);
+        }
+        else {
+            schema = opt.schema;
+        }
+        return new apache_arrow_1.Table(schema, batchesFixed);
     }
-    else {
-        return new apache_arrow_1.Table(columns);
+    const tbl = new apache_arrow_1.Table(columns);
+    if (metadata !== undefined) {
+        // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+        tbl.schema.metadata = metadata;
     }
+    return tbl;
 }
 exports.makeArrowTable = makeArrowTable;
 /**
  * Create an empty Arrow table with the provided schema
  */
-function makeEmptyTable(schema) {
-    return makeArrowTable([], { schema });
+function makeEmptyTable(schema, metadata) {
+    return makeArrowTable([], { schema }, metadata);
 }
 exports.makeEmptyTable = makeEmptyTable;
 /**
@@ -322,9 +463,52 @@ function makeVector(values, type, stringAsDictionary) {
         return (0, apache_arrow_1.vectorFromArray)(values);
     }
 }
+/** Helper function to apply embeddings from metadata to an input table */
+async function applyEmbeddingsFromMetadata(table, schema) {
+    const registry = (0, registry_1.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;
+        const dtype = schema.fields.find((f) => f.name === destColumn).type;
+        if (isFixedSizeList(dtype)) {
+            destType = (0, sanitize_1.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 apache_arrow_1.Table(columns);
+    return alignTable(newTable, schema);
+}
 /** Helper function to apply embeddings to an input table */
 async function applyEmbeddings(table, embeddings, schema) {
-    if (embeddings == null) {
+    if (schema?.metadata.has("embedding_functions")) {
+        return applyEmbeddingsFromMetadata(table, schema);
+    }
+    else if (embeddings == null || embeddings === undefined) {
         return table;
     }
     if (schema !== undefined && schema !== null) {
@@ -339,8 +523,8 @@ async function applyEmbeddings(table, embeddings, schema) {
     });
     const newColumns = Object.fromEntries(colEntries);
     const sourceColumn = newColumns[embeddings.sourceColumn];
-    const destColumn = embeddings.destColumn ?? "vector";
-    const innerDestType = embeddings.embeddingDataType ?? new apache_arrow_1.Float32();
+    const destColumn = embeddings.vectorColumn ?? "vector";
+    const innerDestType = embeddings.function.embeddingDataType() ?? new apache_arrow_1.Float32();
     if (sourceColumn === undefined) {
         throw new Error(`Cannot apply embedding function because the source column '${embeddings.sourceColumn}' was not present in the data`);
     }
@@ -351,8 +535,9 @@ async function applyEmbeddings(table, embeddings, schema) {
             // 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) {
@@ -376,7 +561,7 @@ async function applyEmbeddings(table, embeddings, schema) {
             throw new Error("Internal error: `makeArrowTable` unexpectedly created a table with more than one batch");
         }
         const values = sourceColumn.toArray();
-        const vectors = await embeddings.embed(values);
+        const vectors = await embeddings.function.computeSourceEmbeddings(values);
         if (vectors.length !== values.length) {
             throw new Error("Embedding function did not return an embedding for each input element");
         }
@@ -419,9 +604,10 @@ exports.convertToTable = convertToTable;
 function newVectorType(dim, innerType) {
     // 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 apache_arrow_1.Field("item", innerType, true);
+    const children = new apache_arrow_1.Field("item", (0, sanitize_1.sanitizeType)(innerType), true);
     return new apache_arrow_1.FixedSizeList(dim, children);
 }
+exports.newVectorType = newVectorType;
 /**
  * Serialize an Array of records into a buffer using the Arrow IPC File serialization
  *
@@ -483,12 +669,12 @@ async function fromDataToBuffer(data, embeddings, schema) {
     if (schema !== undefined && schema !== null) {
         schema = (0, sanitize_1.sanitizeSchema)(schema);
     }
-    if (data instanceof apache_arrow_1.Table) {
+    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);
     }
 }
 exports.fromDataToBuffer = fromDataToBuffer;
@@ -548,10 +734,22 @@ function validateSchemaEmbeddings(schema, data, embeddings) {
     // 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 apache_arrow_1.FixedSizeList) {
+    for (let field of schema.fields) {
+        if (isFixedSizeList(field.type)) {
+            field = (0, sanitize_1.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) => f["vectorColumn"] === field.name) ===
+                        undefined) {
+                        missingEmbeddingFields.push(field);
+                    }
+                }
+                else {
+                    missingEmbeddingFields.push(field);
+                }
             }
             else {
                 fields.push(field);

package/nodejs-artifacts/connection.d.ts

@@ -1,4 +1,5 @@
-import { Table as ArrowTable, Schema } from "apache-arrow";
+import { Table as ArrowTable, Schema } from "./arrow";
+import { EmbeddingFunctionConfig } from "./embedding/registry";
 import { ConnectionOptions, Connection as LanceDbConnection } from "./native";
 import { Table } from "./table";
 /**
@@ -39,6 +40,14 @@ export interface CreateTableOptions {
      * The available options are described at https://lancedb.github.io/lancedb/guides/storage/
      */
     storageOptions?: Record<string, string>;
+    /**
+     * If true then data files will be written with the legacy format
+     *
+     * The default is true while the new format is in beta
+     */
+    useLegacyFormat?: boolean;
+    schema?: Schema;
+    embeddingFunction?: EmbeddingFunctionConfig;
 }
 export interface OpenTableOptions {
     /**

package/nodejs-artifacts/connection.js

@@ -14,8 +14,8 @@
 // limitations under the License.
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.Connection = exports.connect = void 0;
-const apache_arrow_1 = require("apache-arrow");
 const arrow_1 = require("./arrow");
+const registry_1 = require("./embedding/registry");
 const native_1 = require("./native");
 const table_1 = require("./table");
 /**
@@ -109,14 +109,14 @@ class Connection {
             mode = "exist_ok";
         }
         let table;
-        if (data instanceof apache_arrow_1.Table) {
+        if ((0, arrow_1.isArrowTable)(data)) {
             table = data;
         }
         else {
-            table = (0, arrow_1.makeArrowTable)(data);
+            table = (0, arrow_1.makeArrowTable)(data, options);
         }
-        const buf = await (0, arrow_1.fromTableToBuffer)(table);
-        const innerTable = await this.inner.createTable(name, buf, mode, cleanseStorageOptions(options?.storageOptions));
+        const buf = await (0, arrow_1.fromTableToBuffer)(table, options?.embeddingFunction, options?.schema);
+        const innerTable = await this.inner.createTable(name, buf, mode, cleanseStorageOptions(options?.storageOptions), options?.useLegacyFormat);
         return new table_1.Table(innerTable);
     }
     /**
@@ -130,9 +130,15 @@ class Connection {
         if (mode === "create" && existOk) {
             mode = "exist_ok";
         }
-        const table = (0, arrow_1.makeEmptyTable)(schema);
+        let metadata = undefined;
+        if (options?.embeddingFunction !== undefined) {
+            const embeddingFunction = options.embeddingFunction;
+            const registry = (0, registry_1.getRegistry)();
+            metadata = registry.getTableMetadata([embeddingFunction]);
+        }
+        const table = (0, arrow_1.makeEmptyTable)(schema, metadata);
         const buf = await (0, arrow_1.fromTableToBuffer)(table);
-        const innerTable = await this.inner.createEmptyTable(name, buf, mode, cleanseStorageOptions(options?.storageOptions));
+        const innerTable = await this.inner.createEmptyTable(name, buf, mode, cleanseStorageOptions(options?.storageOptions), options?.useLegacyFormat);
         return new table_1.Table(innerTable);
     }
     /**

package/nodejs-artifacts/embedding/embedding_function.d.ts

@@ -1,45 +1,71 @@
-import { type Float } from "apache-arrow";
+import "reflect-metadata";
+import { DataType, Float, type IntoVector } from "../arrow";
+/**
+ * Options for a given embedding function
+ */
+export interface FunctionOptions {
+    [key: string]: any;
+}
 /**
  * An embedding function that automatically creates vector representation for a given column.
  */
-export interface EmbeddingFunction<T> {
+export declare abstract class EmbeddingFunction<T = any, M extends FunctionOptions = FunctionOptions> {
     /**
-     * The name of the column that will be used as input for the Embedding Function.
-     */
-    sourceColumn: string;
-    /**
-     * 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
+     *
+     * @param optionsOrDatatype - The options for the field or the datatype
      *
-     * By default this is "vector"
+     * @see {@link lancedb.LanceSchema}
      */
-    destColumn?: string;
+    sourceField(optionsOrDatatype: Partial<FieldOptions> | DataType): [DataType, Map<string, EmbeddingFunction>];
     /**
-     * Should the source column be excluded from the resulting table
+     * vectorField is used in combination with `LanceSchema` to provide a declarative data model
+     *
+     * @param options - The options for the field
      *
-     * By default the source column is included.  Set this to true and
-     * only the embedding will be stored.
+     * @see {@link lancedb.LanceSchema}
      */
-    excludeSource?: boolean;
+    vectorField(optionsOrDatatype?: Partial<FieldOptions> | DataType): [DataType, Map<string, EmbeddingFunction>];
+    /** The number of dimensions of the embeddings */
+    ndims(): number | 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[]>;
+    /**
+    Compute the embeddings for a single query
+   */
+    computeQueryEmbeddings(data: T): Promise<IntoVector>;
+}
+export interface FieldOptions<T extends DataType = DataType> {
+    datatype: T;
+    dims?: number;
 }
-/** Test if the input seems to be an embedding function */
-export declare function isEmbeddingFunction<T>(value: unknown): value is EmbeddingFunction<T>;

package/nodejs-artifacts/embedding/embedding_function.js

@@ -1,5 +1,5 @@
 "use strict";
-// 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.
@@ -13,15 +13,94 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.isEmbeddingFunction = void 0;
-/** Test if the input seems to be an embedding function */
-function isEmbeddingFunction(value) {
-    if (typeof value !== "object" || value === null) {
-        return false;
+exports.EmbeddingFunction = void 0;
+require("reflect-metadata");
+const arrow_1 = require("../arrow");
+const sanitize_1 = require("../sanitize");
+/**
+ * An embedding function that automatically creates vector representation for a given column.
+ */
+class EmbeddingFunction {
+    /**
+     * sourceField is used in combination with `LanceSchema` to provide a declarative data model
+     *
+     * @param optionsOrDatatype - The options for the field or the datatype
+     *
+     * @see {@link lancedb.LanceSchema}
+     */
+    sourceField(optionsOrDatatype) {
+        let datatype = (0, arrow_1.isDataType)(optionsOrDatatype)
+            ? optionsOrDatatype
+            : optionsOrDatatype?.datatype;
+        if (!datatype) {
+            throw new Error("Datatype is required");
+        }
+        datatype = (0, sanitize_1.sanitizeType)(datatype);
+        const metadata = new Map();
+        metadata.set("source_column_for", this);
+        return [datatype, metadata];
     }
-    if (!("sourceColumn" in value) || !("embed" in value)) {
-        return false;
+    /**
+     * vectorField is used in combination with `LanceSchema` to provide a declarative data model
+     *
+     * @param options - The options for the field
+     *
+     * @see {@link lancedb.LanceSchema}
+     */
+    vectorField(optionsOrDatatype) {
+        let dtype;
+        let vectorType;
+        let dims = this.ndims();
+        // `func.vectorField(new Float32())`
+        if ((0, arrow_1.isDataType)(optionsOrDatatype)) {
+            dtype = optionsOrDatatype;
+        }
+        else {
+            // `func.vectorField({
+            //  datatype: new Float32(),
+            //  dims: 10
+            // })`
+            dims = dims ?? optionsOrDatatype?.dims;
+            dtype = optionsOrDatatype?.datatype;
+        }
+        if (dtype !== undefined) {
+            // `func.vectorField(new FixedSizeList(dims, new Field("item", new Float32(), true)))`
+            // or `func.vectorField({datatype: new FixedSizeList(dims, new Field("item", new Float32(), true))})`
+            if ((0, arrow_1.isFixedSizeList)(dtype)) {
+                vectorType = dtype;
+                // `func.vectorField(new Float32())`
+                // or `func.vectorField({datatype: new Float32()})`
+            }
+            else if ((0, arrow_1.isFloat)(dtype)) {
+                // No `ndims` impl and no `{dims: n}` provided;
+                if (dims === undefined) {
+                    throw new Error("ndims is required for vector field");
+                }
+                vectorType = (0, arrow_1.newVectorType)(dims, dtype);
+            }
+            else {
+                throw new Error("Expected FixedSizeList or Float as datatype for vector field");
+            }
+        }
+        else {
+            if (dims === undefined) {
+                throw new Error("ndims is required for vector field");
+            }
+            vectorType = new arrow_1.FixedSizeList(dims, new arrow_1.Field("item", new arrow_1.Float32(), true));
+        }
+        const metadata = new Map();
+        metadata.set("vector_column_for", this);
+        return [vectorType, metadata];
+    }
+    /** The number of dimensions of the embeddings */
+    ndims() {
+        return undefined;
+    }
+    /**
+    Compute the embeddings for a single query
+   */
+    async computeQueryEmbeddings(data) {
+        return this.computeSourceEmbeddings([data]).then((embeddings) => embeddings[0]);
     }
-    return (typeof value.sourceColumn === "string" && typeof value.embed === "function");
 }
-exports.isEmbeddingFunction = isEmbeddingFunction;
+exports.EmbeddingFunction = EmbeddingFunction;

package/nodejs-artifacts/embedding/index.d.ts

@@ -1,2 +1,28 @@
-export { EmbeddingFunction, isEmbeddingFunction } from "./embedding_function";
-export { OpenAIEmbeddingFunction } from "./openai";
+import { Schema } from "../arrow";
+import { EmbeddingFunction } from "./embedding_function";
+export { EmbeddingFunction } from "./embedding_function";
+export * from "./openai";
+export * from "./registry";
+/**
+ * Create a schema with embedding functions.
+ *
+ * @param fields
+ * @returns Schema
+ * @example
+ * ```ts
+ * class MyEmbeddingFunction extends EmbeddingFunction {
+ * // ...
+ * }
+ * const func = new MyEmbeddingFunction();
+ * const schema = LanceSchema({
+ *   id: new Int32(),
+ *   text: func.sourceField(new Utf8()),
+ *   vector: func.vectorField(),
+ *   // optional: specify the datatype and/or dimensions
+ *   vector2: func.vectorField({ datatype: new Float32(), dims: 3}),
+ * });
+ *
+ * const table = await db.createTable("my_table", data, { schema });
+ * ```
+ */
+export declare function LanceSchema(fields: Record<string, [object, Map<string, EmbeddingFunction>] | object>): Schema;