@lancedb/lancedb 0.5.0 → 0.5.2

package/lancedb/query.ts

@@ -12,7 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { Table as ArrowTable, RecordBatch, tableFromIPC } from "apache-arrow";
+import {
+  Table as ArrowTable,
+  type IntoVector,
+  RecordBatch,
+  tableFromIPC,
+} from "./arrow";
 import { type IvfPqOptions } from "./indices";
 import {
   RecordBatchIterator as NativeBatchIterator,
@@ -50,6 +55,39 @@ export class RecordBatchIterator implements AsyncIterator<RecordBatch> {
 }
 /* eslint-enable */
 
+class RecordBatchIterable<
+  NativeQueryType extends NativeQuery | NativeVectorQuery,
+> implements AsyncIterable<RecordBatch>
+{
+  private inner: NativeQueryType;
+  private options?: QueryExecutionOptions;
+
+  constructor(inner: NativeQueryType, options?: QueryExecutionOptions) {
+    this.inner = inner;
+    this.options = options;
+  }
+
+  // biome-ignore lint/suspicious/noExplicitAny: skip
+  [Symbol.asyncIterator](): AsyncIterator<RecordBatch<any>, any, undefined> {
+    return new RecordBatchIterator(
+      this.inner.execute(this.options?.maxBatchLength),
+    );
+  }
+}
+
+/**
+ * Options that control the behavior of a particular query execution
+ */
+export interface QueryExecutionOptions {
+  /**
+   * The maximum number of rows to return in a single batch
+   *
+   * Batches may have fewer rows if the underlying data is stored
+   * in smaller chunks.
+   */
+  maxBatchLength?: number;
+}
+
 /** Common methods supported by all query types */
 export class QueryBase<
   NativeQueryType extends NativeQuery | NativeVectorQuery,
@@ -108,9 +146,12 @@ export class QueryBase<
    * object insertion order is easy to get wrong and `Map` is more foolproof.
    */
   select(
-    columns: string[] | Map<string, string> | Record<string, string>,
+    columns: string[] | Map<string, string> | Record<string, string> | string,
   ): QueryType {
     let columnTuples: [string, string][];
+    if (typeof columns === "string") {
+      columns = [columns];
+    }
     if (Array.isArray(columns)) {
       columnTuples = columns.map((c) => [c, c]);
     } else if (columns instanceof Map) {
@@ -133,8 +174,10 @@ export class QueryBase<
     return this as unknown as QueryType;
   }
 
-  protected nativeExecute(): Promise<NativeBatchIterator> {
-    return this.inner.execute();
+  protected nativeExecute(
+    options?: Partial<QueryExecutionOptions>,
+  ): Promise<NativeBatchIterator> {
+    return this.inner.execute(options?.maxBatchLength);
   }
 
   /**
@@ -148,8 +191,10 @@ export class QueryBase<
    * single query)
    *
    */
-  protected execute(): RecordBatchIterator {
-    return new RecordBatchIterator(this.nativeExecute());
+  protected execute(
+    options?: Partial<QueryExecutionOptions>,
+  ): RecordBatchIterator {
+    return new RecordBatchIterator(this.nativeExecute(options));
   }
 
   // biome-ignore lint/suspicious/noExplicitAny: skip
@@ -159,18 +204,18 @@ export class QueryBase<
   }
 
   /** Collect the results as an Arrow @see {@link ArrowTable}. */
-  async toArrow(): Promise<ArrowTable> {
+  async toArrow(options?: Partial<QueryExecutionOptions>): Promise<ArrowTable> {
     const batches = [];
-    for await (const batch of this) {
+    for await (const batch of new RecordBatchIterable(this.inner, options)) {
       batches.push(batch);
     }
     return new ArrowTable(batches);
   }
 
   /** Collect the results as an array of objects. */
-  async toArray(): Promise<unknown[]> {
-    const tbl = await this.toArrow();
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-return
+  // biome-ignore lint/suspicious/noExplicitAny: arrow.toArrow() returns any[]
+  async toArray(options?: Partial<QueryExecutionOptions>): Promise<any[]> {
+    const tbl = await this.toArrow(options);
     return tbl.toArray();
   }
 }
@@ -369,9 +414,8 @@ export class Query extends QueryBase<NativeQuery, Query> {
    * Vector searches always have a `limit`.  If `limit` has not been called then
    * a default `limit` of 10 will be used.  @see {@link Query#limit}
    */
-  nearestTo(vector: unknown): VectorQuery {
-    // biome-ignore lint/suspicious/noExplicitAny: skip
-    const vectorQuery = this.inner.nearestTo(Float32Array.from(vector as any));
+  nearestTo(vector: IntoVector): VectorQuery {
+    const vectorQuery = this.inner.nearestTo(Float32Array.from(vector));
     return new VectorQuery(vectorQuery);
   }
 }

package/lancedb/sanitize.ts

@@ -20,6 +20,7 @@
 // comes from the exact same library instance.  This is not always the case
 // and so we must sanitize the input to ensure that it is compatible.
 
+import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type";
 import {
   Binary,
   Bool,
@@ -75,10 +76,9 @@ import {
   Uint64,
   Union,
   Utf8,
-} from "apache-arrow";
-import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type";
+} from "./arrow";
 
-function sanitizeMetadata(
+export function sanitizeMetadata(
   metadataLike?: unknown,
 ): Map<string, string> | undefined {
   if (metadataLike === undefined || metadataLike === null) {
@@ -97,7 +97,7 @@ function sanitizeMetadata(
   return metadataLike as Map<string, string>;
 }
 
-function sanitizeInt(typeLike: object) {
+export function sanitizeInt(typeLike: object) {
   if (
     !("bitWidth" in typeLike) ||
     typeof typeLike.bitWidth !== "number" ||
@@ -111,14 +111,14 @@ function sanitizeInt(typeLike: object) {
   return new Int(typeLike.isSigned, typeLike.bitWidth as IntBitWidth);
 }
 
-function sanitizeFloat(typeLike: object) {
+export function sanitizeFloat(typeLike: object) {
   if (!("precision" in typeLike) || typeof typeLike.precision !== "number") {
     throw Error("Expected a Float Type to have a `precision` property");
   }
   return new Float(typeLike.precision as Precision);
 }
 
-function sanitizeDecimal(typeLike: object) {
+export function sanitizeDecimal(typeLike: object) {
   if (
     !("scale" in typeLike) ||
     typeof typeLike.scale !== "number" ||
@@ -134,14 +134,14 @@ function sanitizeDecimal(typeLike: object) {
   return new Decimal(typeLike.scale, typeLike.precision, typeLike.bitWidth);
 }
 
-function sanitizeDate(typeLike: object) {
+export function sanitizeDate(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected a Date type to have a `unit` property");
   }
   return new Date_(typeLike.unit as DateUnit);
 }
 
-function sanitizeTime(typeLike: object) {
+export function sanitizeTime(typeLike: object) {
   if (
     !("unit" in typeLike) ||
     typeof typeLike.unit !== "number" ||
@@ -155,7 +155,7 @@ function sanitizeTime(typeLike: object) {
   return new Time(typeLike.unit, typeLike.bitWidth as TimeBitWidth);
 }
 
-function sanitizeTimestamp(typeLike: object) {
+export function sanitizeTimestamp(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected a Timestamp type to have a `unit` property");
   }
@@ -166,7 +166,7 @@ function sanitizeTimestamp(typeLike: object) {
   return new Timestamp(typeLike.unit, timezone);
 }
 
-function sanitizeTypedTimestamp(
+export function sanitizeTypedTimestamp(
   typeLike: object,
   // eslint-disable-next-line @typescript-eslint/naming-convention
   Datatype:
@@ -182,14 +182,14 @@ function sanitizeTypedTimestamp(
   return new Datatype(timezone);
 }
 
-function sanitizeInterval(typeLike: object) {
+export function sanitizeInterval(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected an Interval type to have a `unit` property");
   }
   return new Interval(typeLike.unit);
 }
 
-function sanitizeList(typeLike: object) {
+export function sanitizeList(typeLike: object) {
   if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
     throw Error(
       "Expected a List type to have an array-like `children` property",
@@ -201,7 +201,7 @@ function sanitizeList(typeLike: object) {
   return new List(sanitizeField(typeLike.children[0]));
 }
 
-function sanitizeStruct(typeLike: object) {
+export function sanitizeStruct(typeLike: object) {
   if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
     throw Error(
       "Expected a Struct type to have an array-like `children` property",
@@ -210,7 +210,7 @@ function sanitizeStruct(typeLike: object) {
   return new Struct(typeLike.children.map((child) => sanitizeField(child)));
 }
 
-function sanitizeUnion(typeLike: object) {
+export function sanitizeUnion(typeLike: object) {
   if (
     !("typeIds" in typeLike) ||
     !("mode" in typeLike) ||
@@ -234,7 +234,7 @@ function sanitizeUnion(typeLike: object) {
   );
 }
 
-function sanitizeTypedUnion(
+export function sanitizeTypedUnion(
   typeLike: object,
   // eslint-disable-next-line @typescript-eslint/naming-convention
   UnionType: typeof DenseUnion | typeof SparseUnion,
@@ -256,7 +256,7 @@ function sanitizeTypedUnion(
   );
 }
 
-function sanitizeFixedSizeBinary(typeLike: object) {
+export function sanitizeFixedSizeBinary(typeLike: object) {
   if (!("byteWidth" in typeLike) || typeof typeLike.byteWidth !== "number") {
     throw Error(
       "Expected a FixedSizeBinary type to have a `byteWidth` property",
@@ -265,7 +265,7 @@ function sanitizeFixedSizeBinary(typeLike: object) {
   return new FixedSizeBinary(typeLike.byteWidth);
 }
 
-function sanitizeFixedSizeList(typeLike: object) {
+export function sanitizeFixedSizeList(typeLike: object) {
   if (!("listSize" in typeLike) || typeof typeLike.listSize !== "number") {
     throw Error("Expected a FixedSizeList type to have a `listSize` property");
   }
@@ -283,7 +283,7 @@ function sanitizeFixedSizeList(typeLike: object) {
   );
 }
 
-function sanitizeMap(typeLike: object) {
+export function sanitizeMap(typeLike: object) {
   if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
     throw Error(
       "Expected a Map type to have an array-like `children` property",
@@ -300,14 +300,14 @@ function sanitizeMap(typeLike: object) {
   );
 }
 
-function sanitizeDuration(typeLike: object) {
+export function sanitizeDuration(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected a Duration type to have a `unit` property");
   }
   return new Duration(typeLike.unit);
 }
 
-function sanitizeDictionary(typeLike: object) {
+export function sanitizeDictionary(typeLike: object) {
   if (!("id" in typeLike) || typeof typeLike.id !== "number") {
     throw Error("Expected a Dictionary type to have an `id` property");
   }
@@ -329,7 +329,7 @@ function sanitizeDictionary(typeLike: object) {
 }
 
 // biome-ignore lint/suspicious/noExplicitAny: skip
-function sanitizeType(typeLike: unknown): DataType<any> {
+export function sanitizeType(typeLike: unknown): DataType<any> {
   if (typeof typeLike !== "object" || typeLike === null) {
     throw Error("Expected a Type but object was null/undefined");
   }
@@ -449,7 +449,7 @@ function sanitizeType(typeLike: unknown): DataType<any> {
   }
 }
 
-function sanitizeField(fieldLike: unknown): Field {
+export function sanitizeField(fieldLike: unknown): Field {
   if (fieldLike instanceof Field) {
     return fieldLike;
   }

package/lancedb/table.ts

@@ -12,8 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { Schema, tableFromIPC } from "apache-arrow";
-import { Data, fromDataToBuffer } from "./arrow";
+import {
+  Table as ArrowTable,
+  Data,
+  IntoVector,
+  Schema,
+  fromDataToBuffer,
+  tableFromIPC,
+} from "./arrow";
+
+import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry";
 import { IndexOptions } from "./indices";
 import {
   AddColumnsSql,
@@ -23,8 +31,8 @@ import {
   Table as _NativeTable,
 } from "./native";
 import { Query, VectorQuery } from "./query";
-
 export { IndexConfig } from "./native";
+
 /**
  * Options for adding data to a table.
  */
@@ -109,6 +117,14 @@ export class Table {
     return this.inner.display();
   }
 
+  async #getEmbeddingFunctions(): Promise<
+    Map<string, EmbeddingFunctionConfig>
+  > {
+    const schema = await this.schema();
+    const registry = getRegistry();
+    return registry.parseFunctions(schema.metadata);
+  }
+
   /** Get the schema of the table. */
   async schema(): Promise<Schema> {
     const schemaBuf = await this.inner.schema();
@@ -122,8 +138,15 @@ export class Table {
    */
   async add(data: Data, options?: Partial<AddDataOptions>): Promise<void> {
     const mode = options?.mode ?? "append";
+    const schema = await this.schema();
+    const registry = getRegistry();
+    const functions = registry.parseFunctions(schema.metadata);
 
-    const buffer = await fromDataToBuffer(data);
+    const buffer = await fromDataToBuffer(
+      data,
+      functions.values().next().value,
+      schema,
+    );
     await this.inner.add(buffer, mode);
   }
 
@@ -263,6 +286,40 @@ export class Table {
     return new Query(this.inner);
   }
 
+  /**
+   * Create a search query to find the nearest neighbors
+   * of the given query vector
+   * @param {string} query - the query. This will be converted to a vector using the table's provided embedding function
+   * @rejects {Error} If no embedding functions are defined in the table
+   */
+  search(query: string): Promise<VectorQuery>;
+  /**
+   * Create a search query to find the nearest neighbors
+   * of the given query vector
+   * @param {IntoVector} query - the query vector
+   */
+  search(query: IntoVector): VectorQuery;
+  search(query: string | IntoVector): Promise<VectorQuery> | VectorQuery {
+    if (typeof query !== "string") {
+      return this.vectorSearch(query);
+    } else {
+      return this.#getEmbeddingFunctions().then(async (functions) => {
+        // TODO: Support multiple embedding functions
+        const embeddingFunc: EmbeddingFunctionConfig | undefined = functions
+          .values()
+          .next().value;
+        if (!embeddingFunc) {
+          return Promise.reject(
+            new Error("No embedding functions are defined in the table"),
+          );
+        }
+        const embeddings =
+          await embeddingFunc.function.computeQueryEmbeddings(query);
+        return this.query().nearestTo(embeddings);
+      });
+    }
+  }
+
   /**
    * Search the table with a given query vector.
    *
@@ -270,7 +327,7 @@ export class Table {
    * is the same thing as calling `nearestTo` on the builder returned
    * by `query`.  @see {@link Query#nearestTo} for more details.
    */
-  vectorSearch(vector: unknown): VectorQuery {
+  vectorSearch(vector: IntoVector): VectorQuery {
     return this.query().nearestTo(vector);
   }
 
@@ -416,4 +473,9 @@ export class Table {
   async listIndices(): Promise<IndexConfig[]> {
     return await this.inner.listIndices();
   }
+
+  /** Return the table as an arrow table */
+  async toArrow(): Promise<ArrowTable> {
+    return await this.query().toArrow();
+  }
 }

package/nodejs-artifacts/arrow.d.ts

@@ -1,6 +1,30 @@
 /// <reference types="node" />
-import { Table as ArrowTable, type Float, Schema } from "apache-arrow";
+import { Table as ArrowTable, Binary, DataType, FixedSizeBinary, FixedSizeList, Float, Int, LargeBinary, List, Null, Schema, Struct, Utf8 } from "apache-arrow";
 import { type EmbeddingFunction } from "./embedding/embedding_function";
+import { EmbeddingFunctionConfig } from "./embedding/registry";
+export * from "apache-arrow";
+export type IntoVector = Float32Array | Float64Array | number[];
+export declare function isArrowTable(value: object): value is ArrowTable;
+export declare function isDataType(value: unknown): value is DataType;
+export declare function isNull(value: unknown): value is Null;
+export declare function isInt(value: unknown): value is Int;
+export declare function isFloat(value: unknown): value is Float;
+export declare function isBinary(value: unknown): value is Binary;
+export declare function isLargeBinary(value: unknown): value is LargeBinary;
+export declare function isUtf8(value: unknown): value is Utf8;
+export declare function isLargeUtf8(value: unknown): value is Utf8;
+export declare function isBool(value: unknown): value is Utf8;
+export declare function isDecimal(value: unknown): value is Utf8;
+export declare function isDate(value: unknown): value is Utf8;
+export declare function isTime(value: unknown): value is Utf8;
+export declare function isTimestamp(value: unknown): value is Utf8;
+export declare function isInterval(value: unknown): value is Utf8;
+export declare function isDuration(value: unknown): value is Utf8;
+export declare function isList(value: unknown): value is List;
+export declare function isStruct(value: unknown): value is Struct;
+export declare function isUnion(value: unknown): value is Struct;
+export declare function isFixedSizeBinary(value: unknown): value is FixedSizeBinary;
+export declare function isFixedSizeList(value: unknown): value is FixedSizeList;
 /** Data type accepted by NodeJS SDK */
 export type Data = Record<string, unknown>[] | ArrowTable;
 export declare class VectorColumnOptions {
@@ -13,6 +37,7 @@ export declare class MakeArrowTableOptions {
     schema?: Schema;
     vectorColumns: Record<string, VectorColumnOptions>;
     embeddings?: EmbeddingFunction<unknown>;
+    embeddingFunction?: EmbeddingFunctionConfig;
     /**
      * If true then string columns will be encoded with dictionary encoding
      *
@@ -117,11 +142,11 @@ export declare class MakeArrowTableOptions {
  * assert.deepEqual(table.schema, schema)
  * ```
  */
-export declare function makeArrowTable(data: Array<Record<string, unknown>>, options?: Partial<MakeArrowTableOptions>): ArrowTable;
+export declare function makeArrowTable(data: Array<Record<string, unknown>>, options?: Partial<MakeArrowTableOptions>, metadata?: Map<string, string>): ArrowTable;
 /**
  * Create an empty Arrow table with the provided schema
  */
-export declare function makeEmptyTable(schema: Schema): ArrowTable;
+export declare function makeEmptyTable(schema: Schema, metadata?: Map<string, string>): ArrowTable;
 /**
  * Convert an Array of records into an Arrow Table, optionally applying an
  * embeddings function to it.
@@ -140,7 +165,9 @@ export declare function makeEmptyTable(schema: Schema): ArrowTable;
  * 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 declare function convertToTable<T>(data: Array<Record<string, unknown>>, embeddings?: EmbeddingFunction<T>, makeTableOptions?: Partial<MakeArrowTableOptions>): Promise<ArrowTable>;
+export declare function convertToTable(data: Array<Record<string, unknown>>, embeddings?: EmbeddingFunctionConfig, makeTableOptions?: Partial<MakeArrowTableOptions>): Promise<ArrowTable>;
+/** Creates the Arrow Type for a Vector column with dimension `dim` */
+export declare function newVectorType<T extends Float>(dim: number, innerType: T): FixedSizeList<T>;
 /**
  * Serialize an Array of records into a buffer using the Arrow IPC File serialization
  *
@@ -148,7 +175,7 @@ export declare function convertToTable<T>(data: Array<Record<string, unknown>>,
  *
  * `schema` is required if data is empty
  */
-export declare function fromRecordsToBuffer<T>(data: Array<Record<string, unknown>>, embeddings?: EmbeddingFunction<T>, schema?: Schema): Promise<Buffer>;
+export declare function fromRecordsToBuffer(data: Array<Record<string, unknown>>, embeddings?: EmbeddingFunctionConfig, schema?: Schema): Promise<Buffer>;
 /**
  * Serialize an Array of records into a buffer using the Arrow IPC Stream serialization
  *
@@ -156,7 +183,7 @@ export declare function fromRecordsToBuffer<T>(data: Array<Record<string, unknow
  *
  * `schema` is required if data is empty
  */
-export declare function fromRecordsToStreamBuffer<T>(data: Array<Record<string, unknown>>, embeddings?: EmbeddingFunction<T>, schema?: Schema): Promise<Buffer>;
+export declare function fromRecordsToStreamBuffer(data: Array<Record<string, unknown>>, embeddings?: EmbeddingFunctionConfig, schema?: Schema): Promise<Buffer>;
 /**
  * Serialize an Arrow Table into a buffer using the Arrow IPC File serialization
  *
@@ -165,7 +192,7 @@ export declare function fromRecordsToStreamBuffer<T>(data: Array<Record<string,
  *
  * `schema` is required if the table is empty
  */
-export declare function fromTableToBuffer<T>(table: ArrowTable, embeddings?: EmbeddingFunction<T>, schema?: Schema): Promise<Buffer>;
+export declare function fromTableToBuffer(table: ArrowTable, embeddings?: EmbeddingFunctionConfig, schema?: Schema): Promise<Buffer>;
 /**
  * Serialize an Arrow Table into a buffer using the Arrow IPC File serialization
  *
@@ -174,7 +201,7 @@ export declare function fromTableToBuffer<T>(table: ArrowTable, embeddings?: Emb
  *
  * `schema` is required if the table is empty
  */
-export declare function fromDataToBuffer<T>(data: Data, embeddings?: EmbeddingFunction<T>, schema?: Schema): Promise<Buffer>;
+export declare function fromDataToBuffer(data: Data, embeddings?: EmbeddingFunctionConfig, schema?: Schema): Promise<Buffer>;
 /**
  * Serialize an Arrow Table into a buffer using the Arrow IPC Stream serialization
  *
@@ -183,7 +210,7 @@ export declare function fromDataToBuffer<T>(data: Data, embeddings?: EmbeddingFu
  *
  * `schema` is required if the table is empty
  */
-export declare function fromTableToStreamBuffer<T>(table: ArrowTable, embeddings?: EmbeddingFunction<T>, schema?: Schema): Promise<Buffer>;
+export declare function fromTableToStreamBuffer(table: ArrowTable, embeddings?: EmbeddingFunctionConfig, schema?: Schema): Promise<Buffer>;
 /**
  * Create an empty table with the given schema
  */