@lancedb/lancedb 0.5.1 → 0.7.0

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);
+  }
+}

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 "./arrow";
+import {
+  Table as ArrowTable,
+  type IntoVector,
+  RecordBatch,
+  tableFromIPC,
+} from "./arrow";
 import { type IvfPqOptions } from "./indices";
 import {
   RecordBatchIterator as NativeBatchIterator,
@@ -50,16 +55,60 @@ export class RecordBatchIterator implements AsyncIterator<RecordBatch> {
 }
 /* eslint-enable */
 
-/** Common methods supported by all query types */
-export class QueryBase<
+class RecordBatchIterable<
   NativeQueryType extends NativeQuery | NativeVectorQuery,
-  QueryType,
 > implements AsyncIterable<RecordBatch>
 {
-  protected constructor(protected inner: NativeQueryType) {
+  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>
+  implements AsyncIterable<RecordBatch>
+{
+  protected constructor(
+    protected inner: NativeQueryType | Promise<NativeQueryType>,
+  ) {
     // intentionally empty
   }
 
+  // call a function on the inner (either a promise or the actual object)
+  protected doCall(fn: (inner: NativeQueryType) => void) {
+    if (this.inner instanceof Promise) {
+      this.inner = this.inner.then((inner) => {
+        fn(inner);
+        return inner;
+      });
+    } else {
+      fn(this.inner);
+    }
+  }
   /**
    * A filter statement to be applied to this query.
    *
@@ -72,9 +121,17 @@ export class QueryBase<
    * Filtering performance can often be improved by creating a scalar index
    * on the filter column(s).
    */
-  where(predicate: string): QueryType {
-    this.inner.onlyIf(predicate);
-    return this as unknown as QueryType;
+  where(predicate: string): this {
+    this.doCall((inner: NativeQueryType) => inner.onlyIf(predicate));
+    return this;
+  }
+  /**
+   * A filter statement to be applied to this query.
+   * @alias where
+   * @deprecated Use `where` instead
+   */
+  filter(predicate: string): this {
+    return this.where(predicate);
   }
 
   /**
@@ -108,9 +165,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>,
-  ): QueryType {
+    columns: string[] | Map<string, string> | Record<string, string> | string,
+  ): this {
     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) {
@@ -118,8 +178,10 @@ export class QueryBase<
     } else {
       columnTuples = Object.entries(columns);
     }
-    this.inner.select(columnTuples);
-    return this as unknown as QueryType;
+    this.doCall((inner: NativeQueryType) => {
+      inner.select(columnTuples);
+    });
+    return this;
   }
 
   /**
@@ -128,13 +190,19 @@ export class QueryBase<
    * By default, a plain search has no limit.  If this method is not
    * called then every valid row from the table will be returned.
    */
-  limit(limit: number): QueryType {
-    this.inner.limit(limit);
-    return this as unknown as QueryType;
+  limit(limit: number): this {
+    this.doCall((inner: NativeQueryType) => inner.limit(limit));
+    return this;
   }
 
-  protected nativeExecute(): Promise<NativeBatchIterator> {
-    return this.inner.execute();
+  protected nativeExecute(
+    options?: Partial<QueryExecutionOptions>,
+  ): Promise<NativeBatchIterator> {
+    if (this.inner instanceof Promise) {
+      return this.inner.then((inner) => inner.execute(options?.maxBatchLength));
+    } else {
+      return this.inner.execute(options?.maxBatchLength);
+    }
   }
 
   /**
@@ -148,8 +216,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,21 +229,48 @@ 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) {
+    let inner;
+    if (this.inner instanceof Promise) {
+      inner = await this.inner;
+    } else {
+      inner = this.inner;
+    }
+    for await (const batch of new RecordBatchIterable(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();
   }
+
+  /**
+   * Generates an explanation of the query execution plan.
+   *
+   * @example
+   * import * as lancedb from "@lancedb/lancedb"
+   * const db = await lancedb.connect("./.lancedb");
+   * const table = await db.createTable("my_table", [
+   *   { vector: [1.1, 0.9], id: "1" },
+   * ]);
+   * const plan = await table.query().nearestTo([0.5, 0.2]).explainPlan();
+   *
+   * @param verbose - If true, provides a more detailed explanation. Defaults to false.
+   * @returns A Promise that resolves to a string containing the query execution plan explanation.
+   */
+  async explainPlan(verbose = false): Promise<string> {
+    if (this.inner instanceof Promise) {
+      return this.inner.then((inner) => inner.explainPlan(verbose));
+    } else {
+      return this.inner.explainPlan(verbose);
+    }
+  }
 }
 
 /**
@@ -188,8 +285,8 @@ export interface ExecutableQuery {}
  *
  * This builder can be reused to execute the query many times.
  */
-export class VectorQuery extends QueryBase<NativeVectorQuery, VectorQuery> {
-  constructor(inner: NativeVectorQuery) {
+export class VectorQuery extends QueryBase<NativeVectorQuery> {
+  constructor(inner: NativeVectorQuery | Promise<NativeVectorQuery>) {
     super(inner);
   }
 
@@ -216,7 +313,8 @@ export class VectorQuery extends QueryBase<NativeVectorQuery, VectorQuery> {
    * you the desired recall.
    */
   nprobes(nprobes: number): VectorQuery {
-    this.inner.nprobes(nprobes);
+    super.doCall((inner) => inner.nprobes(nprobes));
+
     return this;
   }
 
@@ -230,7 +328,7 @@ export class VectorQuery extends QueryBase<NativeVectorQuery, VectorQuery> {
    * whose data type is a fixed-size-list of floats.
    */
   column(column: string): VectorQuery {
-    this.inner.column(column);
+    super.doCall((inner) => inner.column(column));
     return this;
   }
 
@@ -248,8 +346,10 @@ export class VectorQuery extends QueryBase<NativeVectorQuery, VectorQuery> {
    *
    * By default "l2" is used.
    */
-  distanceType(distanceType: string): VectorQuery {
-    this.inner.distanceType(distanceType);
+  distanceType(
+    distanceType: Required<IvfPqOptions>["distanceType"],
+  ): VectorQuery {
+    super.doCall((inner) => inner.distanceType(distanceType));
     return this;
   }
 
@@ -283,7 +383,7 @@ export class VectorQuery extends QueryBase<NativeVectorQuery, VectorQuery> {
    * distance between the query vector and the actual uncompressed vector.
    */
   refineFactor(refineFactor: number): VectorQuery {
-    this.inner.refineFactor(refineFactor);
+    super.doCall((inner) => inner.refineFactor(refineFactor));
     return this;
   }
 
@@ -308,7 +408,7 @@ export class VectorQuery extends QueryBase<NativeVectorQuery, VectorQuery> {
    * factor can often help restore some of the results lost by post filtering.
    */
   postfilter(): VectorQuery {
-    this.inner.postfilter();
+    super.doCall((inner) => inner.postfilter());
     return this;
   }
 
@@ -322,13 +422,13 @@ export class VectorQuery extends QueryBase<NativeVectorQuery, VectorQuery> {
    * calculate your recall to select an appropriate value for nprobes.
    */
   bypassVectorIndex(): VectorQuery {
-    this.inner.bypassVectorIndex();
+    super.doCall((inner) => inner.bypassVectorIndex());
     return this;
   }
 }
 
 /** A builder for LanceDB queries. */
-export class Query extends QueryBase<NativeQuery, Query> {
+export class Query extends QueryBase<NativeQuery> {
   constructor(tbl: NativeTable) {
     super(tbl.query());
   }
@@ -370,9 +470,38 @@ 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));
-    return new VectorQuery(vectorQuery);
+  nearestTo(vector: IntoVector): VectorQuery {
+    if (this.inner instanceof Promise) {
+      const nativeQuery = this.inner.then(async (inner) => {
+        if (vector instanceof Promise) {
+          const arr = await vector.then((v) => Float32Array.from(v));
+          return inner.nearestTo(arr);
+        } else {
+          return inner.nearestTo(Float32Array.from(vector));
+        }
+      });
+      return new VectorQuery(nativeQuery);
+    }
+    if (vector instanceof Promise) {
+      const res = (async () => {
+        try {
+          const v = await vector;
+          const arr = Float32Array.from(v);
+          //
+          // biome-ignore lint/suspicious/noExplicitAny: we need to get the `inner`, but js has no package scoping
+          const value: any = this.nearestTo(arr);
+          const inner = value.inner as
+            | NativeVectorQuery
+            | Promise<NativeVectorQuery>;
+          return inner;
+        } catch (e) {
+          return Promise.reject(e);
+        }
+      })();
+      return new VectorQuery(res);
+    } else {
+      const vectorQuery = this.inner.nearestTo(Float32Array.from(vector));
+      return new VectorQuery(vectorQuery);
+    }
   }
 }

package/lancedb/remote/client.ts

@@ -0,0 +1,221 @@
+// Copyright 2023 LanceDB Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// 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.
+
+import axios, {
+  AxiosError,
+  type AxiosResponse,
+  type ResponseType,
+} from "axios";
+import { Table as ArrowTable } from "../arrow";
+import { tableFromIPC } from "../arrow";
+import { VectorQuery } from "../query";
+
+export class RestfulLanceDBClient {
+  #dbName: string;
+  #region: string;
+  #apiKey: string;
+  #hostOverride?: string;
+  #closed: boolean = false;
+  #connectionTimeout: number = 12 * 1000; // 12 seconds;
+  #readTimeout: number = 30 * 1000; // 30 seconds;
+  #session?: import("axios").AxiosInstance;
+
+  constructor(
+    dbName: string,
+    apiKey: string,
+    region: string,
+    hostOverride?: string,
+    connectionTimeout?: number,
+    readTimeout?: number,
+  ) {
+    this.#dbName = dbName;
+    this.#apiKey = apiKey;
+    this.#region = region;
+    this.#hostOverride = hostOverride ?? this.#hostOverride;
+    this.#connectionTimeout = connectionTimeout ?? this.#connectionTimeout;
+    this.#readTimeout = readTimeout ?? this.#readTimeout;
+  }
+
+  // todo: cache the session.
+  get session(): import("axios").AxiosInstance {
+    if (this.#session !== undefined) {
+      return this.#session;
+    } else {
+      return axios.create({
+        baseURL: this.url,
+        headers: {
+          // biome-ignore lint: external API
+          Authorization: `Bearer ${this.#apiKey}`,
+        },
+        transformResponse: decodeErrorData,
+        timeout: this.#connectionTimeout,
+      });
+    }
+  }
+
+  get url(): string {
+    return (
+      this.#hostOverride ??
+      `https://${this.#dbName}.${this.#region}.api.lancedb.com`
+    );
+  }
+
+  get headers(): { [key: string]: string } {
+    const headers: { [key: string]: string } = {
+      "x-api-key": this.#apiKey,
+      "x-request-id": "na",
+    };
+    if (this.#region == "local") {
+      headers["Host"] = `${this.#dbName}.${this.#region}.api.lancedb.com`;
+    }
+    if (this.#hostOverride) {
+      headers["x-lancedb-database"] = this.#dbName;
+    }
+    return headers;
+  }
+
+  isOpen(): boolean {
+    return !this.#closed;
+  }
+
+  private checkNotClosed(): void {
+    if (this.#closed) {
+      throw new Error("Connection is closed");
+    }
+  }
+
+  close(): void {
+    this.#session = undefined;
+    this.#closed = true;
+  }
+
+  // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+  async get(uri: string, params?: Record<string, any>): Promise<any> {
+    this.checkNotClosed();
+    uri = new URL(uri, this.url).toString();
+    let response;
+    try {
+      response = await this.session.get(uri, {
+        headers: this.headers,
+        params,
+      });
+    } catch (e) {
+      if (e instanceof AxiosError) {
+        response = e.response;
+      } else {
+        throw e;
+      }
+    }
+
+    RestfulLanceDBClient.checkStatus(response!);
+    return response!.data;
+  }
+
+  // biome-ignore lint/suspicious/noExplicitAny: api response
+  async post(uri: string, body?: any): Promise<any>;
+  async post(
+    uri: string,
+    // biome-ignore lint/suspicious/noExplicitAny: api request
+    body: any,
+    additional: {
+      config?: { responseType: "arraybuffer" };
+      headers?: Record<string, string>;
+      params?: Record<string, string>;
+    },
+  ): Promise<Buffer>;
+  async post(
+    uri: string,
+    // biome-ignore lint/suspicious/noExplicitAny: api request
+    body?: any,
+    additional?: {
+      config?: { responseType: ResponseType };
+      headers?: Record<string, string>;
+      params?: Record<string, string>;
+    },
+    // biome-ignore lint/suspicious/noExplicitAny: api response
+  ): Promise<any> {
+    this.checkNotClosed();
+    uri = new URL(uri, this.url).toString();
+    additional = Object.assign(
+      { config: { responseType: "json" } },
+      additional,
+    );
+
+    const headers = { ...this.headers, ...additional.headers };
+
+    if (!headers["Content-Type"]) {
+      headers["Content-Type"] = "application/json";
+    }
+    let response;
+    try {
+      response = await this.session.post(uri, body, {
+        headers,
+        responseType: additional!.config!.responseType,
+        params: new Map(Object.entries(additional.params ?? {})),
+      });
+    } catch (e) {
+      if (e instanceof AxiosError) {
+        response = e.response;
+      } else {
+        throw e;
+      }
+    }
+    RestfulLanceDBClient.checkStatus(response!);
+    if (additional!.config!.responseType === "arraybuffer") {
+      return response!.data;
+    } else {
+      return JSON.parse(response!.data);
+    }
+  }
+
+  async listTables(limit = 10, pageToken = ""): Promise<string[]> {
+    const json = await this.get("/v1/table", { limit, pageToken });
+    return json.tables;
+  }
+
+  async query(tableName: string, query: VectorQuery): Promise<ArrowTable> {
+    const tbl = await this.post(`/v1/table/${tableName}/query`, query, {
+      config: {
+        responseType: "arraybuffer",
+      },
+    });
+    return tableFromIPC(tbl);
+  }
+
+  static checkStatus(response: AxiosResponse): void {
+    if (response.status === 404) {
+      throw new Error(`Not found: ${response.data}`);
+    } else if (response.status >= 400 && response.status < 500) {
+      throw new Error(
+        `Bad Request: ${response.status}, error: ${response.data}`,
+      );
+    } else if (response.status >= 500 && response.status < 600) {
+      throw new Error(
+        `Internal Server Error: ${response.status}, error: ${response.data}`,
+      );
+    } else if (response.status !== 200) {
+      throw new Error(
+        `Unknown Error: ${response.status}, error: ${response.data}`,
+      );
+    }
+  }
+}
+
+function decodeErrorData(data: unknown) {
+  if (Buffer.isBuffer(data)) {
+    const decoded = data.toString("utf-8");
+    return decoded;
+  }
+  return data;
+}