@lancedb/lancedb 0.5.2 → 0.7.1

package/lancedb/query.ts

@@ -89,15 +89,26 @@ export interface QueryExecutionOptions {
 }
 
 /** Common methods supported by all query types */
-export class QueryBase<
-  NativeQueryType extends NativeQuery | NativeVectorQuery,
-  QueryType,
-> implements AsyncIterable<RecordBatch>
+export class QueryBase<NativeQueryType extends NativeQuery | NativeVectorQuery>
+  implements AsyncIterable<RecordBatch>
 {
-  protected constructor(protected inner: NativeQueryType) {
+  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.
    *
@@ -110,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);
   }
 
   /**
@@ -147,7 +166,7 @@ export class QueryBase<
    */
   select(
     columns: string[] | Map<string, string> | Record<string, string> | string,
-  ): QueryType {
+  ): this {
     let columnTuples: [string, string][];
     if (typeof columns === "string") {
       columns = [columns];
@@ -159,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;
   }
 
   /**
@@ -169,15 +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(
     options?: Partial<QueryExecutionOptions>,
   ): Promise<NativeBatchIterator> {
-    return this.inner.execute(options?.maxBatchLength);
+    if (this.inner instanceof Promise) {
+      return this.inner.then((inner) => inner.execute(options?.maxBatchLength));
+    } else {
+      return this.inner.execute(options?.maxBatchLength);
+    }
   }
 
   /**
@@ -206,7 +231,13 @@ export class QueryBase<
   /** Collect the results as an Arrow @see {@link ArrowTable}. */
   async toArrow(options?: Partial<QueryExecutionOptions>): Promise<ArrowTable> {
     const batches = [];
-    for await (const batch of new RecordBatchIterable(this.inner, options)) {
+    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);
@@ -218,6 +249,28 @@ export class QueryBase<
     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);
+    }
+  }
 }
 
 /**
@@ -232,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);
   }
 
@@ -260,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;
   }
 
@@ -274,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;
   }
 
@@ -292,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;
   }
 
@@ -327,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;
   }
 
@@ -352,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;
   }
 
@@ -366,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());
   }
@@ -415,7 +471,37 @@ export class Query extends QueryBase<NativeQuery, Query> {
    * a default `limit` of 10 will be used.  @see {@link Query#limit}
    */
   nearestTo(vector: IntoVector): VectorQuery {
-    const vectorQuery = this.inner.nearestTo(Float32Array.from(vector));
-    return new VectorQuery(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;
+}

package/lancedb/remote/connection.ts

@@ -0,0 +1,201 @@
+import { Schema } from "apache-arrow";
+import {
+  Data,
+  SchemaLike,
+  fromTableToStreamBuffer,
+  makeEmptyTable,
+} from "../arrow";
+import {
+  Connection,
+  CreateTableOptions,
+  OpenTableOptions,
+  TableNamesOptions,
+} from "../connection";
+import { Table } from "../table";
+import { TTLCache } from "../util";
+import { RestfulLanceDBClient } from "./client";
+import { RemoteTable } from "./table";
+
+export interface RemoteConnectionOptions {
+  apiKey?: string;
+  region?: string;
+  hostOverride?: string;
+  connectionTimeout?: number;
+  readTimeout?: number;
+}
+
+export class RemoteConnection extends Connection {
+  #dbName: string;
+  #apiKey: string;
+  #region: string;
+  #client: RestfulLanceDBClient;
+  #tableCache = new TTLCache(300_000);
+
+  constructor(
+    url: string,
+    {
+      apiKey,
+      region,
+      hostOverride,
+      connectionTimeout,
+      readTimeout,
+    }: RemoteConnectionOptions,
+  ) {
+    super();
+    apiKey = apiKey ?? process.env.LANCEDB_API_KEY;
+    region = region ?? process.env.LANCEDB_REGION;
+
+    if (!apiKey) {
+      throw new Error("apiKey is required when connecting to LanceDB Cloud");
+    }
+
+    if (!region) {
+      throw new Error("region is required when connecting to LanceDB Cloud");
+    }
+
+    const parsed = new URL(url);
+    if (parsed.protocol !== "db:") {
+      throw new Error(
+        `invalid protocol: ${parsed.protocol}, only accepts db://`,
+      );
+    }
+
+    this.#dbName = parsed.hostname;
+    this.#apiKey = apiKey;
+    this.#region = region;
+    this.#client = new RestfulLanceDBClient(
+      this.#dbName,
+      this.#apiKey,
+      this.#region,
+      hostOverride,
+      connectionTimeout,
+      readTimeout,
+    );
+  }
+
+  isOpen(): boolean {
+    return this.#client.isOpen();
+  }
+  close(): void {
+    return this.#client.close();
+  }
+
+  display(): string {
+    return `RemoteConnection(${this.#dbName})`;
+  }
+
+  async tableNames(options?: Partial<TableNamesOptions>): Promise<string[]> {
+    const response = await this.#client.get("/v1/table/", {
+      limit: options?.limit ?? 10,
+      // biome-ignore lint/style/useNamingConvention: <explanation>
+      page_token: options?.startAfter ?? "",
+    });
+    const body = await response.body();
+    for (const table of body.tables) {
+      this.#tableCache.set(table, true);
+    }
+    return body.tables;
+  }
+
+  async openTable(
+    name: string,
+    _options?: Partial<OpenTableOptions> | undefined,
+  ): Promise<Table> {
+    if (this.#tableCache.get(name) === undefined) {
+      await this.#client.post(
+        `/v1/table/${encodeURIComponent(name)}/describe/`,
+      );
+      this.#tableCache.set(name, true);
+    }
+    return new RemoteTable(this.#client, name, this.#dbName);
+  }
+
+  async createTable(
+    nameOrOptions:
+      | string
+      | ({ name: string; data: Data } & Partial<CreateTableOptions>),
+    data?: Data,
+    options?: Partial<CreateTableOptions> | undefined,
+  ): Promise<Table> {
+    if (typeof nameOrOptions !== "string" && "name" in nameOrOptions) {
+      const { name, data, ...options } = nameOrOptions;
+      return this.createTable(name, data, options);
+    }
+    if (data === undefined) {
+      throw new Error("data is required");
+    }
+    if (options?.mode) {
+      console.warn(
+        "option 'mode' is not supported in LanceDB Cloud",
+        "LanceDB Cloud only supports the default 'create' mode.",
+        "If the table already exists, an error will be thrown.",
+      );
+    }
+    if (options?.embeddingFunction) {
+      console.warn(
+        "embedding_functions is not yet supported on LanceDB Cloud.",
+        "Please vote https://github.com/lancedb/lancedb/issues/626 ",
+        "for this feature.",
+      );
+    }
+
+    const { buf } = await Table.parseTableData(
+      data,
+      options,
+      true /** streaming */,
+    );
+
+    await this.#client.post(
+      `/v1/table/${encodeURIComponent(nameOrOptions)}/create/`,
+      buf,
+      {
+        config: {
+          responseType: "arraybuffer",
+        },
+        headers: { "Content-Type": "application/vnd.apache.arrow.stream" },
+      },
+    );
+    this.#tableCache.set(nameOrOptions, true);
+    return new RemoteTable(this.#client, nameOrOptions, this.#dbName);
+  }
+
+  async createEmptyTable(
+    name: string,
+    schema: SchemaLike,
+    options?: Partial<CreateTableOptions> | undefined,
+  ): Promise<Table> {
+    if (options?.mode) {
+      console.warn(`mode is not supported on LanceDB Cloud`);
+    }
+
+    if (options?.embeddingFunction) {
+      console.warn(
+        "embeddingFunction is not yet supported on LanceDB Cloud.",
+        "Please vote https://github.com/lancedb/lancedb/issues/626 ",
+        "for this feature.",
+      );
+    }
+    const emptyTable = makeEmptyTable(schema);
+    const buf = await fromTableToStreamBuffer(emptyTable);
+
+    await this.#client.post(
+      `/v1/table/${encodeURIComponent(name)}/create/`,
+      buf,
+      {
+        config: {
+          responseType: "arraybuffer",
+        },
+        headers: { "Content-Type": "application/vnd.apache.arrow.stream" },
+      },
+    );
+
+    this.#tableCache.set(name, true);
+    return new RemoteTable(this.#client, name, this.#dbName);
+  }
+
+  async dropTable(name: string): Promise<void> {
+    await this.#client.post(`/v1/table/${encodeURIComponent(name)}/drop/`);
+
+    this.#tableCache.delete(name);
+  }
+}

package/lancedb/remote/index.ts

@@ -0,0 +1,3 @@
+export { RestfulLanceDBClient } from "./client";
+export { type RemoteConnectionOptions, RemoteConnection } from "./connection";
+export { RemoteTable } from "./table";