@lancedb/lancedb
Version:
LanceDB: A serverless, low-latency vector database for AI applications
253 lines (252 loc) • 10.4 kB
JavaScript
;
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
Object.defineProperty(exports, "__esModule", { value: true });
exports.LocalConnection = exports.Connection = void 0;
exports.cleanseStorageOptions = cleanseStorageOptions;
const arrow_1 = require("./arrow");
const arrow_2 = require("./arrow");
const registry_1 = require("./embedding/registry");
const sanitize_1 = require("./sanitize");
const table_1 = require("./table");
/**
* A LanceDB Connection that allows you to open tables and create new ones.
*
* Connection could be local against filesystem or remote against a server.
*
* A Connection is intended to be a long lived object and may hold open
* resources such as HTTP connection pools. This is generally fine and
* a single connection should be shared if it is going to be used many
* times. However, if you are finished with a connection, you may call
* close to eagerly free these resources. Any call to a Connection
* method after it has been closed will result in an error.
*
* Closing a connection is optional. Connections will automatically
* be closed when they are garbage collected.
*
* Any created tables are independent and will continue to work even if
* the underlying connection has been closed.
* @hideconstructor
*/
class Connection {
[Symbol.for("nodejs.util.inspect.custom")]() {
return this.display();
}
}
exports.Connection = Connection;
/** @hideconstructor */
class LocalConnection extends Connection {
inner;
/** @hidden */
constructor(inner) {
super();
this.inner = inner;
}
isOpen() {
return this.inner.isOpen();
}
close() {
this.inner.close();
}
display() {
return this.inner.display();
}
async tableNames(namespacePathOrOptions, options) {
// Detect if first argument is namespacePath array or options object
let namespacePath;
let tableNamesOptions;
if (Array.isArray(namespacePathOrOptions)) {
// First argument is namespacePath array
namespacePath = namespacePathOrOptions;
tableNamesOptions = options;
}
else {
// First argument is options object (backwards compatibility)
namespacePath = undefined;
tableNamesOptions = namespacePathOrOptions;
}
return this.inner.tableNames(namespacePath ?? [], tableNamesOptions?.startAfter, tableNamesOptions?.limit);
}
async openTable(name, namespacePath, options) {
const innerTable = await this.inner.openTable(name, namespacePath ?? [], cleanseStorageOptions(options?.storageOptions), options?.indexCacheSize);
let table = new table_1.LocalTable(innerTable);
// "main" is the default branch, so treat it as no branch. On a real branch,
// scope and pin in one step (yielding "version V of branch B"); otherwise
// pin the version, if any, against main.
const branch = options?.branch != null && options.branch !== "main"
? options.branch
: undefined;
if (branch != null) {
table = await (await table.branches()).checkout(branch, options?.version);
}
else if (options?.version != null) {
await table.checkout(options.version);
}
return table;
}
async cloneTable(targetTableName, sourceUri, options) {
const innerTable = await this.inner.cloneTable(targetTableName, sourceUri, options?.targetNamespacePath ?? [], options?.sourceVersion ?? null, options?.sourceTag ?? null, options?.isShallow ?? true);
return new table_1.LocalTable(innerTable);
}
getStorageOptions(options) {
if (options?.dataStorageVersion !== undefined) {
if (options.storageOptions === undefined) {
options.storageOptions = {};
}
options.storageOptions["newTableDataStorageVersion"] =
options.dataStorageVersion;
}
if (options?.enableV2ManifestPaths !== undefined) {
if (options.storageOptions === undefined) {
options.storageOptions = {};
}
options.storageOptions["newTableEnableV2ManifestPaths"] =
options.enableV2ManifestPaths ? "true" : "false";
}
return cleanseStorageOptions(options?.storageOptions);
}
async createTable(nameOrOptions, dataOrNamespacePath, namespacePathOrOptions, options) {
if (typeof nameOrOptions !== "string" && "name" in nameOrOptions) {
// First overload: createTable(options, namespacePath?)
const { name, data, ...createOptions } = nameOrOptions;
const namespacePath = dataOrNamespacePath;
return this._createTableImpl(name, data, namespacePath, createOptions);
}
// Second overload: createTable(name, data, namespacePath?, options?)
const name = nameOrOptions;
const data = dataOrNamespacePath;
// Detect if third argument is namespacePath array or options object
let namespacePath;
let createOptions;
if (Array.isArray(namespacePathOrOptions)) {
// Third argument is namespacePath array
namespacePath = namespacePathOrOptions;
createOptions = options;
}
else {
// Third argument is options object (backwards compatibility)
namespacePath = undefined;
createOptions = namespacePathOrOptions;
}
return this._createTableImpl(name, data, namespacePath, createOptions);
}
async _createTableImpl(name, data, namespacePath, options) {
if (data === undefined) {
throw new Error("data is required");
}
const { buf, mode } = await parseTableData(data, options);
const storageOptions = this.getStorageOptions(options);
const innerTable = await this.inner.createTable(name, buf, mode, namespacePath ?? [], storageOptions);
return new table_1.LocalTable(innerTable);
}
async createEmptyTable(name, schema, namespacePathOrOptions, options) {
// Detect if third argument is namespacePath array or options object
let namespacePath;
let createOptions;
if (Array.isArray(namespacePathOrOptions)) {
// Third argument is namespacePath array
namespacePath = namespacePathOrOptions;
createOptions = options;
}
else {
// Third argument is options object (backwards compatibility)
namespacePath = undefined;
createOptions = namespacePathOrOptions;
}
let mode = createOptions?.mode ?? "create";
const existOk = createOptions?.existOk ?? false;
if (mode === "create" && existOk) {
mode = "exist_ok";
}
let metadata = undefined;
if (createOptions?.embeddingFunction !== undefined) {
const embeddingFunction = createOptions.embeddingFunction;
const registry = (0, registry_1.getRegistry)();
metadata = registry.getTableMetadata([embeddingFunction]);
}
const storageOptions = this.getStorageOptions(createOptions);
const table = (0, arrow_2.makeEmptyTable)(schema, metadata);
const buf = await (0, arrow_2.fromTableToBuffer)(table);
const innerTable = await this.inner.createEmptyTable(name, buf, mode, namespacePath ?? [], storageOptions);
return new table_1.LocalTable(innerTable);
}
async dropTable(name, namespacePath) {
return this.inner.dropTable(name, namespacePath ?? []);
}
async dropAllTables(namespacePath) {
return this.inner.dropAllTables(namespacePath ?? []);
}
describeNamespace(namespacePath) {
return this.inner.describeNamespace(namespacePath);
}
listNamespaces(namespacePath, options) {
return this.inner.listNamespaces(namespacePath ?? [], options?.pageToken, options?.limit);
}
createNamespace(namespacePath, options) {
return this.inner.createNamespace(namespacePath, options?.mode, options?.properties);
}
dropNamespace(namespacePath, options) {
return this.inner.dropNamespace(namespacePath, options?.mode, options?.behavior);
}
async renameTable(currentName, newName, options) {
return this.inner.renameTable(currentName, newName, options?.namespacePath ?? [], options?.newNamespacePath);
}
}
exports.LocalConnection = LocalConnection;
/**
* Takes storage options and makes all the keys snake case.
*/
function cleanseStorageOptions(options) {
if (options === undefined) {
return undefined;
}
const result = {};
for (const [key, value] of Object.entries(options)) {
if (value !== undefined) {
const newKey = camelToSnakeCase(key);
result[newKey] = value;
}
}
return result;
}
/**
* Convert a string to snake case. It might already be snake case, in which case it is
* returned unchanged.
*/
function camelToSnakeCase(camel) {
if (camel.includes("_")) {
// Assume if there is at least one underscore, it is already snake case
return camel;
}
if (camel.toLocaleUpperCase() === camel) {
// Assume if the string is all uppercase, it is already snake case
return camel;
}
let result = camel.replace(/[A-Z]/g, (letter) => `_${letter.toLowerCase()}`);
if (result.startsWith("_")) {
result = result.slice(1);
}
return result;
}
async function parseTableData(data, options, streaming = false) {
let mode = options?.mode ?? "create";
const existOk = options?.existOk ?? false;
if (mode === "create" && existOk) {
mode = "exist_ok";
}
let table;
if ((0, arrow_1.isArrowTable)(data)) {
table = (0, sanitize_1.sanitizeTable)(data);
}
else {
table = (0, arrow_1.makeArrowTable)(data, options);
}
if (streaming) {
const buf = await (0, arrow_1.fromTableToStreamBuffer)(table, options?.embeddingFunction, options?.schema);
return { buf, mode };
}
else {
const buf = await (0, arrow_2.fromTableToBuffer)(table, options?.embeddingFunction, options?.schema);
return { buf, mode };
}
}