@lancedb/lancedb 0.5.2 → 0.7.0

package/lancedb/connection.ts

@@ -12,38 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { Table as ArrowTable, Schema } from "./arrow";
-import {
-  fromTableToBuffer,
-  isArrowTable,
-  makeArrowTable,
-  makeEmptyTable,
-} from "./arrow";
+import { Data, Schema, SchemaLike, TableLike } from "./arrow";
+import { fromTableToBuffer, makeEmptyTable } from "./arrow";
 import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry";
-import { ConnectionOptions, Connection as LanceDbConnection } from "./native";
-import { Table } from "./table";
-
-/**
- * Connect to a LanceDB instance at the given URI.
- *
- * Accepted formats:
- *
- * - `/path/to/database` - local database
- * - `s3://bucket/path/to/database` or `gs://bucket/path/to/database` - database on cloud storage
- * - `db://host:port` - remote database (LanceDB cloud)
- * @param {string} uri - The uri of the database. If the database uri starts
- * with `db://` then it connects to a remote database.
- * @see {@link ConnectionOptions} for more details on the URI format.
- */
-export async function connect(
-  uri: string,
-  opts?: Partial<ConnectionOptions>,
-): Promise<Connection> {
-  opts = opts ?? {};
-  opts.storageOptions = cleanseStorageOptions(opts.storageOptions);
-  const nativeConn = await LanceDbConnection.new(uri, opts);
-  return new Connection(nativeConn);
-}
+import { Connection as LanceDbConnection } from "./native";
+import { LocalTable, Table } from "./table";
 
 export interface CreateTableOptions {
   /**
@@ -77,7 +50,7 @@ export interface CreateTableOptions {
    * The default is true while the new format is in beta
    */
   useLegacyFormat?: boolean;
-  schema?: Schema;
+  schema?: SchemaLike;
   embeddingFunction?: EmbeddingFunctionConfig;
 }
 
@@ -117,7 +90,6 @@ export interface TableNamesOptions {
   /** An optional limit to the number of results to return. */
   limit?: number;
 }
-
 /**
  * A LanceDB Connection that allows you to open tables and create new ones.
  *
@@ -136,17 +108,15 @@ export interface TableNamesOptions {
  * Any created tables are independent and will continue to work even if
  * the underlying connection has been closed.
  */
-export class Connection {
-  readonly inner: LanceDbConnection;
-
-  constructor(inner: LanceDbConnection) {
-    this.inner = inner;
+export abstract class Connection {
+  [Symbol.for("nodejs.util.inspect.custom")](): string {
+    return this.display();
   }
 
-  /** Return true if the connection has not been closed */
-  isOpen(): boolean {
-    return this.inner.isOpen();
-  }
+  /**
+   * Return true if the connection has not been closed
+   */
+  abstract isOpen(): boolean;
 
   /**
    * Close the connection, releasing any underlying resources.
@@ -155,14 +125,12 @@ export class Connection {
    *
    * Any attempt to use the connection after it is closed will result in an error.
    */
-  close(): void {
-    this.inner.close();
-  }
+  abstract close(): void;
 
-  /** Return a brief description of the connection */
-  display(): string {
-    return this.inner.display();
-  }
+  /**
+   * Return a brief description of the connection
+   */
+  abstract display(): string;
 
   /**
    * List all the table names in this database.
@@ -170,15 +138,86 @@ export class Connection {
    * Tables will be returned in lexicographical order.
    * @param {Partial<TableNamesOptions>} options - options to control the
    * paging / start point
+   *
    */
-  async tableNames(options?: Partial<TableNamesOptions>): Promise<string[]> {
-    return this.inner.tableNames(options?.startAfter, options?.limit);
-  }
+  abstract tableNames(options?: Partial<TableNamesOptions>): Promise<string[]>;
 
   /**
    * Open a table in the database.
    * @param {string} name - The name of the table
    */
+  abstract openTable(
+    name: string,
+    options?: Partial<OpenTableOptions>,
+  ): Promise<Table>;
+
+  /**
+   * Creates a new Table and initialize it with new data.
+   * @param {object} options - The options object.
+   * @param {string} options.name - The name of the table.
+   * @param {Data} options.data - Non-empty Array of Records to be inserted into the table
+   *
+   */
+  abstract createTable(
+    options: {
+      name: string;
+      data: Data;
+    } & Partial<CreateTableOptions>,
+  ): Promise<Table>;
+  /**
+   * Creates a new Table and initialize it with new data.
+   * @param {string} name - The name of the table.
+   * @param {Record<string, unknown>[] | TableLike} data - Non-empty Array of Records
+   * to be inserted into the table
+   */
+  abstract createTable(
+    name: string,
+    data: Record<string, unknown>[] | TableLike,
+    options?: Partial<CreateTableOptions>,
+  ): Promise<Table>;
+
+  /**
+   * Creates a new empty Table
+   * @param {string} name - The name of the table.
+   * @param {Schema} schema - The schema of the table
+   */
+  abstract createEmptyTable(
+    name: string,
+    schema: import("./arrow").SchemaLike,
+    options?: Partial<CreateTableOptions>,
+  ): Promise<Table>;
+
+  /**
+   * Drop an existing table.
+   * @param {string} name The name of the table to drop.
+   */
+  abstract dropTable(name: string): Promise<void>;
+}
+
+export class LocalConnection extends Connection {
+  readonly inner: LanceDbConnection;
+
+  constructor(inner: LanceDbConnection) {
+    super();
+    this.inner = inner;
+  }
+
+  isOpen(): boolean {
+    return this.inner.isOpen();
+  }
+
+  close(): void {
+    this.inner.close();
+  }
+
+  display(): string {
+    return this.inner.display();
+  }
+
+  async tableNames(options?: Partial<TableNamesOptions>): Promise<string[]> {
+    return this.inner.tableNames(options?.startAfter, options?.limit);
+  }
+
   async openTable(
     name: string,
     options?: Partial<OpenTableOptions>,
@@ -189,58 +228,38 @@ export class Connection {
       options?.indexCacheSize,
     );
 
-    return new Table(innerTable);
+    return new LocalTable(innerTable);
   }
 
-  /**
-   * Creates a new Table and initialize it with new data.
-   * @param {string} name - The name of the table.
-   * @param {Record<string, unknown>[] | ArrowTable} data - Non-empty Array of Records
-   * to be inserted into the table
-   */
   async createTable(
-    name: string,
-    data: Record<string, unknown>[] | ArrowTable,
+    nameOrOptions:
+      | string
+      | ({ name: string; data: Data } & Partial<CreateTableOptions>),
+    data?: Record<string, unknown>[] | TableLike,
     options?: Partial<CreateTableOptions>,
   ): Promise<Table> {
-    let mode: string = options?.mode ?? "create";
-    const existOk = options?.existOk ?? false;
-
-    if (mode === "create" && existOk) {
-      mode = "exist_ok";
+    if (typeof nameOrOptions !== "string" && "name" in nameOrOptions) {
+      const { name, data, ...options } = nameOrOptions;
+      return this.createTable(name, data, options);
     }
-
-    let table: ArrowTable;
-    if (isArrowTable(data)) {
-      table = data;
-    } else {
-      table = makeArrowTable(data, options);
+    if (data === undefined) {
+      throw new Error("data is required");
     }
-
-    const buf = await fromTableToBuffer(
-      table,
-      options?.embeddingFunction,
-      options?.schema,
-    );
+    const { buf, mode } = await Table.parseTableData(data, options);
     const innerTable = await this.inner.createTable(
-      name,
+      nameOrOptions,
       buf,
       mode,
       cleanseStorageOptions(options?.storageOptions),
       options?.useLegacyFormat,
     );
 
-    return new Table(innerTable);
+    return new LocalTable(innerTable);
   }
 
-  /**
-   * Creates a new empty Table
-   * @param {string} name - The name of the table.
-   * @param {Schema} schema - The schema of the table
-   */
   async createEmptyTable(
     name: string,
-    schema: Schema,
+    schema: import("./arrow").SchemaLike,
     options?: Partial<CreateTableOptions>,
   ): Promise<Table> {
     let mode: string = options?.mode ?? "create";
@@ -265,13 +284,9 @@ export class Connection {
       cleanseStorageOptions(options?.storageOptions),
       options?.useLegacyFormat,
     );
-    return new Table(innerTable);
+    return new LocalTable(innerTable);
   }
 
-  /**
-   * Drop an existing table.
-   * @param {string} name The name of the table to drop.
-   */
   async dropTable(name: string): Promise<void> {
     return this.inner.dropTable(name);
   }
@@ -280,7 +295,7 @@ export class Connection {
 /**
  * Takes storage options and makes all the keys snake case.
  */
-function cleanseStorageOptions(
+export function cleanseStorageOptions(
   options?: Record<string, string>,
 ): Record<string, string> | undefined {
   if (options === undefined) {

package/lancedb/embedding/embedding_function.ts

@@ -35,6 +35,11 @@ export interface FunctionOptions {
   [key: string]: any;
 }
 
+export interface EmbeddingFunctionConstructor<
+  T extends EmbeddingFunction = EmbeddingFunction,
+> {
+  new (modelOptions?: T["TOptions"]): T;
+}
 /**
  * An embedding function that automatically creates vector representation for a given column.
  */
@@ -43,6 +48,12 @@ export abstract class EmbeddingFunction<
   T = any,
   M extends FunctionOptions = FunctionOptions,
 > {
+  /**
+   * @ignore
+   *  This is only used for associating the options type with the class for type checking
+   */
+  // biome-ignore lint/style/useNamingConvention: we want to keep the name as it is
+  readonly TOptions!: M;
   /**
    * Convert the embedding function to a JSON object
    * It is used to serialize the embedding function to the schema
@@ -170,7 +181,7 @@ export abstract class EmbeddingFunction<
   /**
   Compute the embeddings for a single query
  */
-  async computeQueryEmbeddings(data: T): Promise<IntoVector> {
+  async computeQueryEmbeddings(data: T): Promise<Awaited<IntoVector>> {
     return this.computeSourceEmbeddings([data]).then(
       (embeddings) => embeddings[0],
     );

package/lancedb/embedding/openai.ts

@@ -13,24 +13,29 @@
 // limitations under the License.
 
 import type OpenAI from "openai";
+import { type EmbeddingCreateParams } from "openai/resources";
 import { Float, Float32 } from "../arrow";
 import { EmbeddingFunction } from "./embedding_function";
 import { register } from "./registry";
 
 export type OpenAIOptions = {
-  apiKey?: string;
-  model?: string;
+  apiKey: string;
+  model: EmbeddingCreateParams["model"];
 };
 
 @register("openai")
 export class OpenAIEmbeddingFunction extends EmbeddingFunction<
   string,
-  OpenAIOptions
+  Partial<OpenAIOptions>
 > {
   #openai: OpenAI;
-  #modelName: string;
+  #modelName: OpenAIOptions["model"];
 
-  constructor(options: OpenAIOptions = { model: "text-embedding-ada-002" }) {
+  constructor(
+    options: Partial<OpenAIOptions> = {
+      model: "text-embedding-ada-002",
+    },
+  ) {
     super();
     const openAIKey = options?.apiKey ?? process.env.OPENAI_API_KEY;
     if (!openAIKey) {
@@ -73,7 +78,7 @@ export class OpenAIEmbeddingFunction extends EmbeddingFunction<
       case "text-embedding-3-small":
         return 1536;
       default:
-        return null as never;
+        throw new Error(`Unknown model: ${this.#modelName}`);
     }
   }
 

package/lancedb/embedding/registry.ts

@@ -12,21 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import type { EmbeddingFunction } from "./embedding_function";
+import {
+  type EmbeddingFunction,
+  type EmbeddingFunctionConstructor,
+} from "./embedding_function";
 import "reflect-metadata";
-
-export interface EmbeddingFunctionOptions {
-  [key: string]: unknown;
-}
-
-export interface EmbeddingFunctionFactory<
-  T extends EmbeddingFunction = EmbeddingFunction,
-> {
-  new (modelOptions?: EmbeddingFunctionOptions): T;
-}
+import { OpenAIEmbeddingFunction } from "./openai";
 
 interface EmbeddingFunctionCreate<T extends EmbeddingFunction> {
-  create(options?: EmbeddingFunctionOptions): T;
+  create(options?: T["TOptions"]): T;
 }
 
 /**
@@ -36,7 +30,7 @@ interface EmbeddingFunctionCreate<T extends EmbeddingFunction> {
  * or TextEmbeddingFunction and registering it with the registry
  */
 export class EmbeddingFunctionRegistry {
-  #functions: Map<string, EmbeddingFunctionFactory> = new Map();
+  #functions = new Map<string, EmbeddingFunctionConstructor>();
 
   /**
    * Register an embedding function
@@ -44,7 +38,9 @@ export class EmbeddingFunctionRegistry {
    * @param func The function to register
    * @throws Error if the function is already registered
    */
-  register<T extends EmbeddingFunctionFactory = EmbeddingFunctionFactory>(
+  register<
+    T extends EmbeddingFunctionConstructor = EmbeddingFunctionConstructor,
+  >(
     this: EmbeddingFunctionRegistry,
     alias?: string,
     // biome-ignore lint/suspicious/noExplicitAny: <explanation>
@@ -69,18 +65,34 @@ export class EmbeddingFunctionRegistry {
    * Fetch an embedding function by name
    * @param name The name of the function
    */
-  get<T extends EmbeddingFunction<unknown> = EmbeddingFunction>(
-    name: string,
-  ): EmbeddingFunctionCreate<T> | undefined {
+  get<T extends EmbeddingFunction<unknown>, Name extends string = "">(
+    name: Name extends "openai" ? "openai" : string,
+    //This makes it so that you can use string constants as "types", or use an explicitly supplied type
+    // ex:
+    // `registry.get("openai") -> EmbeddingFunctionCreate<OpenAIEmbeddingFunction>`
+    // `registry.get<MyCustomEmbeddingFunction>("my_func") -> EmbeddingFunctionCreate<MyCustomEmbeddingFunction> | undefined`
+    //
+    // the reason this is important is that we always know our built in functions are defined so the user isnt forced to do a non null/undefined
+    // ```ts
+    // const openai: OpenAIEmbeddingFunction = registry.get("openai").create()
+    // ```
+  ): Name extends "openai"
+    ? EmbeddingFunctionCreate<OpenAIEmbeddingFunction>
+    : EmbeddingFunctionCreate<T> | undefined {
+    type Output = Name extends "openai"
+      ? EmbeddingFunctionCreate<OpenAIEmbeddingFunction>
+      : EmbeddingFunctionCreate<T> | undefined;
+
     const factory = this.#functions.get(name);
     if (!factory) {
-      return undefined;
+      return undefined as Output;
     }
+
     return {
-      create: function (options: EmbeddingFunctionOptions) {
-        return new factory(options) as unknown as T;
+      create: function (options?: T["TOptions"]) {
+        return new factory(options);
       },
-    };
+    } as Output;
   }
 
   /**
@@ -104,7 +116,7 @@ export class EmbeddingFunctionRegistry {
         name: string;
         sourceColumn: string;
         vectorColumn: string;
-        model: EmbeddingFunctionOptions;
+        model: EmbeddingFunction["TOptions"];
       };
       const functions = <FunctionConfig[]>(
         JSON.parse(metadata.get("embedding_functions")!)

package/lancedb/index.ts

@@ -12,25 +12,43 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import {
+  Connection,
+  LocalConnection,
+  cleanseStorageOptions,
+} from "./connection";
+
+import {
+  ConnectionOptions,
+  Connection as LanceDbConnection,
+} from "./native.js";
+
+import { RemoteConnection, RemoteConnectionOptions } from "./remote";
+
 export {
   WriteOptions,
   WriteMode,
   AddColumnsSql,
   ColumnAlteration,
   ConnectionOptions,
+  IndexStatistics,
+  IndexMetadata,
+  IndexConfig,
 } from "./native.js";
+
 export {
   makeArrowTable,
   MakeArrowTableOptions,
   Data,
   VectorColumnOptions,
 } from "./arrow";
+
 export {
-  connect,
   Connection,
   CreateTableOptions,
   TableNamesOptions,
 } from "./connection";
+
 export {
   ExecutableQuery,
   Query,
@@ -38,6 +56,87 @@ export {
   VectorQuery,
   RecordBatchIterator,
 } from "./query";
+
 export { Index, IndexOptions, IvfPqOptions } from "./indices";
-export { Table, AddDataOptions, IndexConfig, UpdateOptions } from "./table";
+
+export { Table, AddDataOptions, UpdateOptions } from "./table";
+
 export * as embedding from "./embedding";
+
+/**
+ * Connect to a LanceDB instance at the given URI.
+ *
+ * Accepted formats:
+ *
+ * - `/path/to/database` - local database
+ * - `s3://bucket/path/to/database` or `gs://bucket/path/to/database` - database on cloud storage
+ * - `db://host:port` - remote database (LanceDB cloud)
+ * @param {string} uri - The uri of the database. If the database uri starts
+ * with `db://` then it connects to a remote database.
+ * @see {@link ConnectionOptions} for more details on the URI format.
+ * @example
+ * ```ts
+ * const conn = await connect("/path/to/database");
+ * ```
+ * @example
+ * ```ts
+ * const conn = await connect(
+ *   "s3://bucket/path/to/database",
+ *   {storageOptions: {timeout: "60s"}
+ * });
+ * ```
+ */
+export async function connect(
+  uri: string,
+  opts?: Partial<ConnectionOptions | RemoteConnectionOptions>,
+): Promise<Connection>;
+/**
+ * Connect to a LanceDB instance at the given URI.
+ *
+ * Accepted formats:
+ *
+ * - `/path/to/database` - local database
+ * - `s3://bucket/path/to/database` or `gs://bucket/path/to/database` - database on cloud storage
+ * - `db://host:port` - remote database (LanceDB cloud)
+ * @param  options - The options to use when connecting to the database
+ * @see {@link ConnectionOptions} for more details on the URI format.
+ * @example
+ * ```ts
+ * const conn = await connect({
+ *   uri: "/path/to/database",
+ *   storageOptions: {timeout: "60s"}
+ * });
+ * ```
+ */
+export async function connect(
+  opts: Partial<RemoteConnectionOptions | ConnectionOptions> & { uri: string },
+): Promise<Connection>;
+export async function connect(
+  uriOrOptions:
+    | string
+    | (Partial<RemoteConnectionOptions | ConnectionOptions> & { uri: string }),
+  opts: Partial<ConnectionOptions | RemoteConnectionOptions> = {},
+): Promise<Connection> {
+  let uri: string | undefined;
+  if (typeof uriOrOptions !== "string") {
+    const { uri: uri_, ...options } = uriOrOptions;
+    uri = uri_;
+    opts = options;
+  } else {
+    uri = uriOrOptions;
+  }
+
+  if (!uri) {
+    throw new Error("uri is required");
+  }
+
+  if (uri?.startsWith("db://")) {
+    return new RemoteConnection(uri, opts as RemoteConnectionOptions);
+  }
+  opts = (opts as ConnectionOptions) ?? {};
+  (<ConnectionOptions>opts).storageOptions = cleanseStorageOptions(
+    (<ConnectionOptions>opts).storageOptions,
+  );
+  const nativeConn = await LanceDbConnection.new(uri, opts);
+  return new LocalConnection(nativeConn);
+}

package/lancedb/merge.ts

@@ -0,0 +1,70 @@
+import { Data, fromDataToBuffer } from "./arrow";
+import { NativeMergeInsertBuilder } from "./native";
+
+/** A builder used to create and run a merge insert operation */
+export class MergeInsertBuilder {
+  #native: NativeMergeInsertBuilder;
+
+  /** Construct a MergeInsertBuilder. __Internal use only.__ */
+  constructor(native: NativeMergeInsertBuilder) {
+    this.#native = native;
+  }
+
+  /**
+   * Rows that exist in both the source table (new data) and
+   * the target table (old data) will be updated, replacing
+   * the old row with the corresponding matching row.
+   *
+   * If there are multiple matches then the behavior is undefined.
+   * Currently this causes multiple copies of the row to be created
+   * but that behavior is subject to change.
+   *
+   * An optional condition may be specified.  If it is, then only
+   * matched rows that satisfy the condtion will be updated.  Any
+   * rows that do not satisfy the condition will be left as they
+   * are.  Failing to satisfy the condition does not cause a
+   * "matched row" to become a "not matched" row.
+   *
+   * The condition should be an SQL string.  Use the prefix
+   * target. to refer to rows in the target table (old data)
+   * and the prefix source. to refer to rows in the source
+   * table (new data).
+   *
+   * For example, "target.last_update < source.last_update"
+   */
+  whenMatchedUpdateAll(options?: { where: string }): MergeInsertBuilder {
+    return new MergeInsertBuilder(
+      this.#native.whenMatchedUpdateAll(options?.where),
+    );
+  }
+  /**
+   * Rows that exist only in the source table (new data) should
+   * be inserted into the target table.
+   */
+  whenNotMatchedInsertAll(): MergeInsertBuilder {
+    return new MergeInsertBuilder(this.#native.whenNotMatchedInsertAll());
+  }
+  /**
+   * Rows that exist only in the target table (old data) will be
+   * deleted.  An optional condition can be provided to limit what
+   * data is deleted.
+   *
+   * @param options.where - An optional condition to limit what data is deleted
+   */
+  whenNotMatchedBySourceDelete(options?: {
+    where: string;
+  }): MergeInsertBuilder {
+    return new MergeInsertBuilder(
+      this.#native.whenNotMatchedBySourceDelete(options?.where),
+    );
+  }
+  /**
+   * Executes the merge insert operation
+   *
+   * Nothing is returned but the `Table` is updated
+   */
+  async execute(data: Data): Promise<void> {
+    const buffer = await fromDataToBuffer(data);
+    await this.#native.execute(buffer);
+  }
+}