@lancedb/lancedb
Version:
LanceDB: A serverless, low-latency vector database for AI applications
201 lines (200 loc) • 9.15 kB
JavaScript
;
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
Object.defineProperty(exports, "__esModule", { value: true });
exports.Scannable = void 0;
const arrow_1 = require("./arrow");
const native_js_1 = require("./native.js");
/**
* A data source that can be scanned as a stream of Arrow `RecordBatch`es.
*
* `Scannable` wraps the schema + optional row count + rescannable flag and
* a callback that yields batches one at a time. It is passed to consumers
* (e.g. `Table.add`, `createTable`, `mergeInsert` — follow-up work) that
* need to pull data without materializing the full dataset in JS memory.
*
* Batches cross the JS↔Rust boundary as Arrow IPC Stream messages; a fresh
* writer serializes each batch, and the Rust side decodes it with
* `arrow_ipc::reader::StreamReader`. One batch is in flight at a time.
*/
class Scannable {
schema;
numRows;
rescannable;
/** @hidden */
native;
constructor(native, schema, numRows, rescannable) {
this.native = native;
this.schema = schema;
this.numRows = numRows;
this.rescannable = rescannable;
}
/** @hidden Access the native handle for passing through to Rust consumers. */
get inner() {
return this.native;
}
/**
* Build a Scannable from an explicit schema and a factory that returns a
* fresh batch iterator on each call.
*
* The factory is invoked once per scan. Each iterator yields
* `RecordBatch`es matching the declared schema. Use this when you need
* direct control over the pull loop — for example, to wrap a streaming
* source whose batches are produced lazily.
*
* @param schema - The Arrow schema of the produced batches.
* @param factory - Called at the start of each scan to produce a batch
* iterator. Must be idempotent when `rescannable` is true.
* @param opts - Optional hints. `rescannable` defaults to `true`; set to
* `false` if calling `factory()` twice would not reproduce the same data.
*/
static async fromFactory(schema, factory, opts = {}) {
const numRows = opts.numRows ?? null;
if (numRows != null && !Number.isInteger(numRows)) {
throw new TypeError("numRows must be an integer");
}
const rescannable = opts.rescannable ?? true;
let iter = null;
const getNextBatch = async (isStart) => {
// `isStart` is true on the first pull of every new scan_as_stream.
// Drop any cached iterator so factory() is re-invoked for the next scan
if (isStart) {
iter = null;
}
if (iter === null) {
iter = normalizeIterator(factory());
}
const result = await iter.next();
if (result.done) {
iter = null;
return null;
}
return (0, arrow_1.fromRecordBatchToStreamBuffer)(result.value);
};
const schemaBuf = await (0, arrow_1.fromTableToBuffer)((0, arrow_1.makeEmptyTable)(schema));
const native = new native_js_1.NapiScannable(schemaBuf, numRows, rescannable, getNextBatch);
return new Scannable(native, schema, numRows, rescannable);
}
/**
* Build a Scannable from an in-memory Arrow `Table`. Always rescannable;
* the table's batches are replayed on each scan.
*
* The table's row count is authoritative: `opts.numRows` must either be
* omitted or equal to `table.numRows`. `opts.rescannable` of `false` is
* rejected because in-memory Tables are always rescannable.
*/
static async fromTable(table, opts = {}) {
if (opts.numRows != null && opts.numRows !== table.numRows) {
throw new TypeError(`opts.numRows (${opts.numRows}) does not match table.numRows (${table.numRows}). ` +
`The table's row count is authoritative; omit numRows or pass the matching value.`);
}
if (opts.rescannable === false) {
throw new TypeError(`fromTable does not accept rescannable: false. ` +
`In-memory Arrow Tables are always rescannable; omit the option or pass true.`);
}
return Scannable.fromFactory(table.schema, () => table.batches, {
numRows: table.numRows,
rescannable: true,
});
}
/**
* Build a Scannable from an iterable of `RecordBatch`es. `rescannable`
* defaults to `false`. Pass an explicit schema so the consumer can
* validate before any batch is pulled.
*
* `opts.rescannable: true` is honest for replayable iterables (Arrays,
* Sets, or custom iterables whose `[Symbol.iterator]()` returns a fresh
* iterator each call). It is rejected for one-shot iterables (generators,
* async generators, or already-an-iterator inputs) because their
* `[Symbol.iterator]()` returns the same exhausted object on the second
* scan. For replayable sources outside this shape, use
* `fromFactory(schema, () => createIter(), { rescannable: true })`.
*
* Note: when `opts.rescannable` is `true`, the constructor calls
* `[Symbol.iterator]()` once on the input to perform the structural check.
*/
static async fromIterable(schema, iter, opts = {}) {
if (opts.rescannable === true && isOneShotIterable(iter)) {
throw new TypeError(`fromIterable: rescannable: true is not honest for one-shot iterables ` +
`(generators, async generators, or iterators where [Symbol.iterator]() ` +
`returns the same object). The source would be exhausted after the first scan. ` +
`Use fromFactory(schema, () => createIter(), { rescannable: true }) for sources ` +
`where each call mints a fresh iterator.`);
}
return Scannable.fromFactory(schema, () => iter, {
numRows: opts.numRows,
rescannable: opts.rescannable ?? false,
});
}
/**
* Build a Scannable from an Arrow `RecordBatchReader`. A reader can only
* be consumed once; `rescannable` defaults to `false`.
*
* The reader must already be opened (via `.open()`) so its `.schema` is
* populated. `RecordBatchReader.from(...)` returns an unopened reader.
*
* `opts.rescannable: true` is rejected because `RecordBatchReader` is a
* self-iterator (its `[Symbol.iterator]()` returns itself), and this
* constructor does not call `reader.reset()` between scans, so a second
* scan would always see an exhausted reader. For genuinely replayable
* sources, use
* `fromFactory(schema, () => openReader(), { rescannable: true })`,
* which mints a fresh reader on each scan.
*/
static async fromRecordBatchReader(reader, opts = {}) {
if (opts.rescannable === true) {
throw new TypeError(`fromRecordBatchReader does not accept rescannable: true. ` +
`RecordBatchReader is a self-iterator (its [Symbol.iterator]() ` +
`returns itself) and would be exhausted after the first scan. ` +
`Use fromFactory(schema, () => openReader(), { rescannable: true }) ` +
`for sources where each call mints a fresh reader.`);
}
return Scannable.fromFactory(reader.schema, () => reader, {
numRows: opts.numRows,
rescannable: false,
});
}
}
exports.Scannable = Scannable;
function normalizeIterator(source) {
if (source == null) {
throw new TypeError("Scannable factory returned null/undefined");
}
if (typeof source[Symbol.asyncIterator] === "function") {
return source[Symbol.asyncIterator]();
}
if (typeof source[Symbol.iterator] === "function") {
return source[Symbol.iterator]();
}
// Already an iterator (has `.next`).
if (typeof source.next === "function") {
return source;
}
throw new TypeError("Scannable factory returned a non-iterable value");
}
// A "self-iterator" returns the same object from `[Symbol.iterator]()` /
// `[Symbol.asyncIterator]()`. Generators behave this way, so they exhaust
// after one pass. Replayable iterables (Array, Set, custom) return a fresh
// iterator each call. Detection mirrors `normalizeIterator`'s ordering so
// classification matches scan-time behavior.
function isOneShotIterable(source) {
// null/undefined are not one-shot in any meaningful sense; let
// `normalizeIterator` raise the actual error at scan time.
if (source == null)
return false;
const ref = source;
if (typeof source[Symbol.asyncIterator] ===
"function") {
const it = source[Symbol.asyncIterator]();
return it === ref;
}
if (typeof source[Symbol.iterator] === "function") {
const it = source[Symbol.iterator]();
return it === ref;
}
// Already-an-iterator (has `.next` but no `Symbol.iterator`) is by
// definition one-shot.
if (typeof source.next === "function")
return true;
return false;
}