@lancedb/lancedb 0.5.1 → 0.7.0

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";

package/lancedb/remote/table.ts

@@ -0,0 +1,226 @@
+// 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 { Table as ArrowTable } from "apache-arrow";
+
+import { Data, IntoVector } from "../arrow";
+
+import { IndexStatistics } from "..";
+import { CreateTableOptions } from "../connection";
+import { IndexOptions } from "../indices";
+import { MergeInsertBuilder } from "../merge";
+import { VectorQuery } from "../query";
+import { AddDataOptions, Table, UpdateOptions } from "../table";
+import { IntoSql, toSQL } from "../util";
+import { RestfulLanceDBClient } from "./client";
+
+export class RemoteTable extends Table {
+  #client: RestfulLanceDBClient;
+  #name: string;
+
+  // Used in the display() method
+  #dbName: string;
+
+  get #tablePrefix() {
+    return `/v1/table/${encodeURIComponent(this.#name)}/`;
+  }
+
+  get name(): string {
+    return this.#name;
+  }
+
+  public constructor(
+    client: RestfulLanceDBClient,
+    tableName: string,
+    dbName: string,
+  ) {
+    super();
+    this.#client = client;
+    this.#name = tableName;
+    this.#dbName = dbName;
+  }
+
+  isOpen(): boolean {
+    return !this.#client.isOpen();
+  }
+
+  close(): void {
+    this.#client.close();
+  }
+
+  display(): string {
+    return `RemoteTable(${this.#dbName}; ${this.#name})`;
+  }
+
+  async schema(): Promise<import("apache-arrow").Schema> {
+    const resp = await this.#client.post(`${this.#tablePrefix}/describe/`);
+    // TODO: parse this into a valid arrow schema
+    return resp.schema;
+  }
+  async add(data: Data, options?: Partial<AddDataOptions>): Promise<void> {
+    const { buf, mode } = await Table.parseTableData(
+      data,
+      options as CreateTableOptions,
+      true,
+    );
+    await this.#client.post(`${this.#tablePrefix}/insert/`, buf, {
+      params: {
+        mode,
+      },
+      headers: {
+        "Content-Type": "application/vnd.apache.arrow.stream",
+      },
+    });
+  }
+
+  async update(
+    optsOrUpdates:
+      | (Map<string, string> | Record<string, string>)
+      | ({
+          values: Map<string, IntoSql> | Record<string, IntoSql>;
+        } & Partial<UpdateOptions>)
+      | ({
+          valuesSql: Map<string, string> | Record<string, string>;
+        } & Partial<UpdateOptions>),
+    options?: Partial<UpdateOptions>,
+  ): Promise<void> {
+    const isValues =
+      "values" in optsOrUpdates && typeof optsOrUpdates.values !== "string";
+    const isValuesSql =
+      "valuesSql" in optsOrUpdates &&
+      typeof optsOrUpdates.valuesSql !== "string";
+    const isMap = (obj: unknown): obj is Map<string, string> => {
+      return obj instanceof Map;
+    };
+
+    let predicate;
+    let columns: [string, string][];
+    switch (true) {
+      case isMap(optsOrUpdates):
+        columns = Array.from(optsOrUpdates.entries());
+        predicate = options?.where;
+        break;
+      case isValues && isMap(optsOrUpdates.values):
+        columns = Array.from(optsOrUpdates.values.entries()).map(([k, v]) => [
+          k,
+          toSQL(v),
+        ]);
+        predicate = optsOrUpdates.where;
+        break;
+      case isValues && !isMap(optsOrUpdates.values):
+        columns = Object.entries(optsOrUpdates.values).map(([k, v]) => [
+          k,
+          toSQL(v),
+        ]);
+        predicate = optsOrUpdates.where;
+        break;
+
+      case isValuesSql && isMap(optsOrUpdates.valuesSql):
+        columns = Array.from(optsOrUpdates.valuesSql.entries());
+        predicate = optsOrUpdates.where;
+        break;
+      case isValuesSql && !isMap(optsOrUpdates.valuesSql):
+        columns = Object.entries(optsOrUpdates.valuesSql).map(([k, v]) => [
+          k,
+          v,
+        ]);
+        predicate = optsOrUpdates.where;
+        break;
+      default:
+        columns = Object.entries(optsOrUpdates as Record<string, string>);
+        predicate = options?.where;
+    }
+
+    await this.#client.post(`${this.#tablePrefix}/update/`, {
+      predicate: predicate ?? null,
+      updates: columns,
+    });
+  }
+  async countRows(filter?: unknown): Promise<number> {
+    const payload = { predicate: filter };
+    return await this.#client.post(`${this.#tablePrefix}/count_rows/`, payload);
+  }
+
+  async delete(predicate: unknown): Promise<void> {
+    const payload = { predicate };
+    await this.#client.post(`${this.#tablePrefix}/delete/`, payload);
+  }
+  async createIndex(
+    column: string,
+    options?: Partial<IndexOptions>,
+  ): Promise<void> {
+    if (options !== undefined) {
+      console.warn("options are not yet supported on the LanceDB cloud");
+    }
+    const indexType = "vector";
+    const metric = "L2";
+    const data = {
+      column,
+      // biome-ignore lint/style/useNamingConvention: external API
+      index_type: indexType,
+      // biome-ignore lint/style/useNamingConvention: external API
+      metric_type: metric,
+    };
+    await this.#client.post(`${this.#tablePrefix}/create_index`, data);
+  }
+  query(): import("..").Query {
+    throw new Error("query() is not yet supported on the LanceDB cloud");
+  }
+
+  search(_query: string | IntoVector): VectorQuery {
+    throw new Error("search() is not yet supported on the LanceDB cloud");
+  }
+  vectorSearch(_vector: unknown): import("..").VectorQuery {
+    throw new Error("vectorSearch() is not yet supported on the LanceDB cloud");
+  }
+  addColumns(_newColumnTransforms: unknown): Promise<void> {
+    throw new Error("addColumns() is not yet supported on the LanceDB cloud");
+  }
+  alterColumns(_columnAlterations: unknown): Promise<void> {
+    throw new Error("alterColumns() is not yet supported on the LanceDB cloud");
+  }
+  dropColumns(_columnNames: unknown): Promise<void> {
+    throw new Error("dropColumns() is not yet supported on the LanceDB cloud");
+  }
+  async version(): Promise<number> {
+    const resp = await this.#client.post(`${this.#tablePrefix}/describe/`);
+    return resp.version;
+  }
+  checkout(_version: unknown): Promise<void> {
+    throw new Error("checkout() is not yet supported on the LanceDB cloud");
+  }
+  checkoutLatest(): Promise<void> {
+    throw new Error(
+      "checkoutLatest() is not yet supported on the LanceDB cloud",
+    );
+  }
+  restore(): Promise<void> {
+    throw new Error("restore() is not yet supported on the LanceDB cloud");
+  }
+  optimize(_options?: unknown): Promise<import("../native").OptimizeStats> {
+    throw new Error("optimize() is not yet supported on the LanceDB cloud");
+  }
+  async listIndices(): Promise<import("../native").IndexConfig[]> {
+    return await this.#client.post(`${this.#tablePrefix}/index/list/`);
+  }
+  toArrow(): Promise<ArrowTable> {
+    throw new Error("toArrow() is not yet supported on the LanceDB cloud");
+  }
+  mergeInsert(_on: string | string[]): MergeInsertBuilder {
+    throw new Error("mergeInsert() is not yet supported on the LanceDB cloud");
+  }
+  async indexStats(_name: string): Promise<IndexStatistics | undefined> {
+    throw new Error("indexStats() is not yet supported on the LanceDB cloud");
+  }
+}

package/lancedb/sanitize.ts

@@ -20,10 +20,12 @@
 // 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 { BufferType, Data } from "apache-arrow";
 import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type";
 import {
   Binary,
   Bool,
+  DataLike,
   DataType,
   DateDay,
   DateMillisecond,
@@ -56,9 +58,14 @@ import {
   Map_,
   Null,
   type Precision,
+  RecordBatch,
+  RecordBatchLike,
   Schema,
+  SchemaLike,
   SparseUnion,
   Struct,
+  Table,
+  TableLike,
   Time,
   TimeMicrosecond,
   TimeMillisecond,
@@ -488,7 +495,7 @@ export function sanitizeField(fieldLike: unknown): Field {
  * instance because they might be using a different instance of apache-arrow
  * than lancedb is using.
  */
-export function sanitizeSchema(schemaLike: unknown): Schema {
+export function sanitizeSchema(schemaLike: SchemaLike): Schema {
   if (schemaLike instanceof Schema) {
     return schemaLike;
   }
@@ -514,3 +521,68 @@ export function sanitizeSchema(schemaLike: unknown): Schema {
   );
   return new Schema(sanitizedFields, metadata);
 }
+
+export function sanitizeTable(tableLike: TableLike): Table {
+  if (tableLike instanceof Table) {
+    return tableLike;
+  }
+  if (typeof tableLike !== "object" || tableLike === null) {
+    throw Error("Expected a Table but object was null/undefined");
+  }
+  if (!("schema" in tableLike)) {
+    throw Error(
+      "The table passed in does not appear to be a table (no 'schema' property)",
+    );
+  }
+  if (!("batches" in tableLike)) {
+    throw Error(
+      "The table passed in does not appear to be a table (no 'columns' property)",
+    );
+  }
+  const schema = sanitizeSchema(tableLike.schema);
+
+  const batches = tableLike.batches.map(sanitizeRecordBatch);
+  return new Table(schema, batches);
+}
+
+function sanitizeRecordBatch(batchLike: RecordBatchLike): RecordBatch {
+  if (batchLike instanceof RecordBatch) {
+    return batchLike;
+  }
+  if (typeof batchLike !== "object" || batchLike === null) {
+    throw Error("Expected a RecordBatch but object was null/undefined");
+  }
+  if (!("schema" in batchLike)) {
+    throw Error(
+      "The record batch passed in does not appear to be a record batch (no 'schema' property)",
+    );
+  }
+  if (!("data" in batchLike)) {
+    throw Error(
+      "The record batch passed in does not appear to be a record batch (no 'data' property)",
+    );
+  }
+  const schema = sanitizeSchema(batchLike.schema);
+  const data = sanitizeData(batchLike.data);
+  return new RecordBatch(schema, data);
+}
+function sanitizeData(
+  dataLike: DataLike,
+  // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+): import("apache-arrow").Data<Struct<any>> {
+  if (dataLike instanceof Data) {
+    return dataLike;
+  }
+  return new Data(
+    dataLike.type,
+    dataLike.offset,
+    dataLike.length,
+    dataLike.nullCount,
+    {
+      [BufferType.OFFSET]: dataLike.valueOffsets,
+      [BufferType.DATA]: dataLike.values,
+      [BufferType.VALIDITY]: dataLike.nullBitmap,
+      [BufferType.TYPE]: dataLike.typeIds,
+    },
+  );
+}