@lancedb/lancedb 0.5.2 → 0.7.1

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