mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-10 15:22:25 +00:00
feat(nodejs): feature parity [6/N] - make public interface work with multiple arrow versions (#1392)
previously we didnt have great compatibility with other versions of apache arrow. This should bridge that gap a bit. depends on https://github.com/lancedb/lancedb/pull/1391 see actual diff here https://github.com/universalmind303/lancedb/compare/query-filter...universalmind303:arrow-compatibility
This commit is contained in:
+74
-21
@@ -15,6 +15,7 @@
|
||||
import {
|
||||
Table as ArrowTable,
|
||||
Binary,
|
||||
BufferType,
|
||||
DataType,
|
||||
Field,
|
||||
FixedSizeBinary,
|
||||
@@ -37,14 +38,68 @@ import {
|
||||
type makeTable,
|
||||
vectorFromArray,
|
||||
} from "apache-arrow";
|
||||
import { Buffers } from "apache-arrow/data";
|
||||
import { type EmbeddingFunction } from "./embedding/embedding_function";
|
||||
import { EmbeddingFunctionConfig, getRegistry } from "./embedding/registry";
|
||||
import { sanitizeField, sanitizeSchema, sanitizeType } from "./sanitize";
|
||||
import {
|
||||
sanitizeField,
|
||||
sanitizeSchema,
|
||||
sanitizeTable,
|
||||
sanitizeType,
|
||||
} from "./sanitize";
|
||||
export * from "apache-arrow";
|
||||
export type SchemaLike =
|
||||
| Schema
|
||||
| {
|
||||
fields: FieldLike[];
|
||||
metadata: Map<string, string>;
|
||||
get names(): unknown[];
|
||||
};
|
||||
export type FieldLike =
|
||||
| Field
|
||||
| {
|
||||
type: string;
|
||||
name: string;
|
||||
nullable?: boolean;
|
||||
metadata?: Map<string, string>;
|
||||
};
|
||||
|
||||
export type DataLike =
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
| import("apache-arrow").Data<Struct<any>>
|
||||
| {
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
type: any;
|
||||
length: number;
|
||||
offset: number;
|
||||
stride: number;
|
||||
nullable: boolean;
|
||||
children: DataLike[];
|
||||
get nullCount(): number;
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
values: Buffers<any>[BufferType.DATA];
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
typeIds: Buffers<any>[BufferType.TYPE];
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
nullBitmap: Buffers<any>[BufferType.VALIDITY];
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
valueOffsets: Buffers<any>[BufferType.OFFSET];
|
||||
};
|
||||
|
||||
export type RecordBatchLike =
|
||||
| RecordBatch
|
||||
| {
|
||||
schema: SchemaLike;
|
||||
data: DataLike;
|
||||
};
|
||||
|
||||
export type TableLike =
|
||||
| ArrowTable
|
||||
| { schema: SchemaLike; batches: RecordBatchLike[] };
|
||||
|
||||
export type IntoVector = Float32Array | Float64Array | number[];
|
||||
|
||||
export function isArrowTable(value: object): value is ArrowTable {
|
||||
export function isArrowTable(value: object): value is TableLike {
|
||||
if (value instanceof ArrowTable) return true;
|
||||
return "schema" in value && "batches" in value;
|
||||
}
|
||||
@@ -135,7 +190,7 @@ export function isFixedSizeList(value: unknown): value is FixedSizeList {
|
||||
}
|
||||
|
||||
/** Data type accepted by NodeJS SDK */
|
||||
export type Data = Record<string, unknown>[] | ArrowTable;
|
||||
export type Data = Record<string, unknown>[] | TableLike;
|
||||
|
||||
/*
|
||||
* Options to control how a column should be converted to a vector array
|
||||
@@ -162,7 +217,7 @@ export class MakeArrowTableOptions {
|
||||
* The schema must be specified if there are no records (e.g. to make
|
||||
* an empty table)
|
||||
*/
|
||||
schema?: Schema;
|
||||
schema?: SchemaLike;
|
||||
|
||||
/*
|
||||
* Mapping from vector column name to expected type
|
||||
@@ -310,7 +365,7 @@ export function makeArrowTable(
|
||||
if (opt.schema !== undefined && opt.schema !== null) {
|
||||
opt.schema = sanitizeSchema(opt.schema);
|
||||
opt.schema = validateSchemaEmbeddings(
|
||||
opt.schema,
|
||||
opt.schema as Schema,
|
||||
data,
|
||||
options?.embeddingFunction,
|
||||
);
|
||||
@@ -394,7 +449,7 @@ export function makeArrowTable(
|
||||
// `new ArrowTable(schema, batches)` which does not do any schema inference
|
||||
const firstTable = new ArrowTable(columns);
|
||||
const batchesFixed = firstTable.batches.map(
|
||||
(batch) => new RecordBatch(opt.schema!, batch.data),
|
||||
(batch) => new RecordBatch(opt.schema as Schema, batch.data),
|
||||
);
|
||||
let schema: Schema;
|
||||
if (metadata !== undefined) {
|
||||
@@ -407,9 +462,9 @@ export function makeArrowTable(
|
||||
}
|
||||
}
|
||||
|
||||
schema = new Schema(opt.schema.fields, schemaMetadata);
|
||||
schema = new Schema(opt.schema.fields as Field[], schemaMetadata);
|
||||
} else {
|
||||
schema = opt.schema;
|
||||
schema = opt.schema as Schema;
|
||||
}
|
||||
return new ArrowTable(schema, batchesFixed);
|
||||
}
|
||||
@@ -425,7 +480,7 @@ export function makeArrowTable(
|
||||
* Create an empty Arrow table with the provided schema
|
||||
*/
|
||||
export function makeEmptyTable(
|
||||
schema: Schema,
|
||||
schema: SchemaLike,
|
||||
metadata?: Map<string, string>,
|
||||
): ArrowTable {
|
||||
return makeArrowTable([], { schema }, metadata);
|
||||
@@ -563,17 +618,16 @@ async function applyEmbeddingsFromMetadata(
|
||||
async function applyEmbeddings<T>(
|
||||
table: ArrowTable,
|
||||
embeddings?: EmbeddingFunctionConfig,
|
||||
schema?: Schema,
|
||||
schema?: SchemaLike,
|
||||
): Promise<ArrowTable> {
|
||||
if (schema?.metadata.has("embedding_functions")) {
|
||||
return applyEmbeddingsFromMetadata(table, schema!);
|
||||
} else if (embeddings == null || embeddings === undefined) {
|
||||
return table;
|
||||
}
|
||||
|
||||
if (schema !== undefined && schema !== null) {
|
||||
schema = sanitizeSchema(schema);
|
||||
}
|
||||
if (schema?.metadata.has("embedding_functions")) {
|
||||
return applyEmbeddingsFromMetadata(table, schema! as Schema);
|
||||
} else if (embeddings == null || embeddings === undefined) {
|
||||
return table;
|
||||
}
|
||||
|
||||
// Convert from ArrowTable to Record<String, Vector>
|
||||
const colEntries = [...Array(table.numCols).keys()].map((_, idx) => {
|
||||
@@ -650,7 +704,7 @@ async function applyEmbeddings<T>(
|
||||
`When using embedding functions and specifying a schema the schema should include the embedding column but the column ${destColumn} was missing`,
|
||||
);
|
||||
}
|
||||
return alignTable(newTable, schema);
|
||||
return alignTable(newTable, schema as Schema);
|
||||
}
|
||||
return newTable;
|
||||
}
|
||||
@@ -744,7 +798,7 @@ export async function fromRecordsToStreamBuffer(
|
||||
export async function fromTableToBuffer(
|
||||
table: ArrowTable,
|
||||
embeddings?: EmbeddingFunctionConfig,
|
||||
schema?: Schema,
|
||||
schema?: SchemaLike,
|
||||
): Promise<Buffer> {
|
||||
if (schema !== undefined && schema !== null) {
|
||||
schema = sanitizeSchema(schema);
|
||||
@@ -771,7 +825,7 @@ export async function fromDataToBuffer(
|
||||
schema = sanitizeSchema(schema);
|
||||
}
|
||||
if (isArrowTable(data)) {
|
||||
return fromTableToBuffer(data, embeddings, schema);
|
||||
return fromTableToBuffer(sanitizeTable(data), embeddings, schema);
|
||||
} else {
|
||||
const table = await convertToTable(data, embeddings, { schema });
|
||||
return fromTableToBuffer(table);
|
||||
@@ -789,7 +843,7 @@ export async function fromDataToBuffer(
|
||||
export async function fromTableToStreamBuffer(
|
||||
table: ArrowTable,
|
||||
embeddings?: EmbeddingFunctionConfig,
|
||||
schema?: Schema,
|
||||
schema?: SchemaLike,
|
||||
): Promise<Buffer> {
|
||||
const tableWithEmbeddings = await applyEmbeddings(table, embeddings, schema);
|
||||
const writer = RecordBatchStreamWriter.writeAll(tableWithEmbeddings);
|
||||
@@ -854,7 +908,6 @@ function validateSchemaEmbeddings(
|
||||
for (let field of schema.fields) {
|
||||
if (isFixedSizeList(field.type)) {
|
||||
field = sanitizeField(field);
|
||||
|
||||
if (data.length !== 0 && data?.[0]?.[field.name] === undefined) {
|
||||
if (schema.metadata.has("embedding_functions")) {
|
||||
const embeddings = JSON.parse(
|
||||
|
||||
Reference in New Issue
Block a user