mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-26 05:26:14 +00:00
fix: update scoring autoprojection bindings for current main
Merge main, reuse its existing query request flag and scanner handling, and remove the unused TypeScript error suppression. Cover local projections, deferred queries, and remote option forwarding.
This commit is contained in:
+1545
-911
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,39 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
// Every Apache Arrow release the package accepts as a peer dependency, as
|
||||
// separately installed copies so the tests can feed objects built by one
|
||||
// Arrow version into a LanceDB running on another. When a new Arrow major
|
||||
// is released, add the alias to `package.json` and extend this list; the
|
||||
// test files pick the change up from here.
|
||||
import * as arrow15 from "apache-arrow-15";
|
||||
import * as arrow16 from "apache-arrow-16";
|
||||
import * as arrow17 from "apache-arrow-17";
|
||||
import * as arrow18 from "apache-arrow-18";
|
||||
import * as arrow19 from "apache-arrow-19";
|
||||
import * as arrow20 from "apache-arrow-20";
|
||||
import * as arrow21 from "apache-arrow-21";
|
||||
|
||||
export { arrow15, arrow16, arrow17, arrow18, arrow19, arrow20, arrow21 };
|
||||
|
||||
export type ApacheArrow =
|
||||
| typeof arrow15
|
||||
| typeof arrow16
|
||||
| typeof arrow17
|
||||
| typeof arrow18
|
||||
| typeof arrow19
|
||||
| typeof arrow20
|
||||
| typeof arrow21;
|
||||
|
||||
export const arrowVersions: ApacheArrow[] = [
|
||||
arrow15,
|
||||
arrow16,
|
||||
arrow17,
|
||||
arrow18,
|
||||
arrow19,
|
||||
arrow20,
|
||||
arrow21,
|
||||
];
|
||||
|
||||
/** The newest supported Arrow release, for tests that only need one. */
|
||||
export const latestArrow = arrow21;
|
||||
@@ -0,0 +1,695 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import {
|
||||
Field,
|
||||
Float,
|
||||
Float32,
|
||||
Int32,
|
||||
Int64,
|
||||
List,
|
||||
Schema,
|
||||
Struct,
|
||||
Utf8,
|
||||
} from "apache-arrow";
|
||||
import {
|
||||
convertToTable,
|
||||
makeArrowTable,
|
||||
resolveBlobInputs,
|
||||
} from "../lancedb/arrow";
|
||||
import { BlobFile, blob, coerceBlobValue, isBlobField } from "../lancedb/blob";
|
||||
import { EmbeddingFunction, LanceSchema } from "../lancedb/embedding";
|
||||
import { getRegistry } from "../lancedb/embedding/registry";
|
||||
|
||||
function legacyBlob(name: string): Field {
|
||||
const field = blob(name);
|
||||
return new Field(
|
||||
name,
|
||||
new Struct((field.type as Struct).children.slice(0, 2)),
|
||||
field.nullable,
|
||||
field.metadata,
|
||||
);
|
||||
}
|
||||
|
||||
describe("blob()", () => {
|
||||
it("declares the four-field range layout", () => {
|
||||
const field = blob("image", {
|
||||
nullable: false,
|
||||
inlineSizeThreshold: 12,
|
||||
});
|
||||
expect(field.nullable).toBe(false);
|
||||
expect(isBlobField(field)).toBe(true);
|
||||
expect(
|
||||
field.metadata.get("lance-encoding:blob-inline-size-threshold"),
|
||||
).toBe("12");
|
||||
expect(
|
||||
(field.type as Struct).children.map((f) => [
|
||||
f.name,
|
||||
f.type.toString(),
|
||||
f.nullable,
|
||||
]),
|
||||
).toEqual([
|
||||
["data", "LargeBinary", true],
|
||||
["uri", "Utf8", true],
|
||||
["position", "Uint64", true],
|
||||
["size", "Uint64", true],
|
||||
]);
|
||||
});
|
||||
|
||||
it("marks the field as lance.blob.v2", () => {
|
||||
const field = blob("image", { nullable: false });
|
||||
expect(field.nullable).toBe(false);
|
||||
expect(isBlobField(field)).toBe(true);
|
||||
expect(field.metadata.get("ARROW:extension:name")).toBe("lance.blob.v2");
|
||||
});
|
||||
|
||||
it("writes encoding thresholds as field metadata", () => {
|
||||
const field = blob("video", {
|
||||
inlineSizeThreshold: 1024,
|
||||
dedicatedSizeThreshold: 2 * 1024 * 1024,
|
||||
packFileSizeThreshold: 64 * 1024 * 1024,
|
||||
});
|
||||
expect(
|
||||
field.metadata.get("lance-encoding:blob-inline-size-threshold"),
|
||||
).toBe("1024");
|
||||
expect(
|
||||
field.metadata.get("lance-encoding:blob-dedicated-size-threshold"),
|
||||
).toBe(String(2 * 1024 * 1024));
|
||||
expect(
|
||||
field.metadata.get("lance-encoding:blob-pack-file-size-threshold"),
|
||||
).toBe(String(64 * 1024 * 1024));
|
||||
});
|
||||
|
||||
it("rejects invalid thresholds", () => {
|
||||
expect(() => blob("image", { inlineSizeThreshold: -1 })).toThrow(
|
||||
/inlineSizeThreshold must be non-negative/,
|
||||
);
|
||||
expect(() => blob("image", { dedicatedSizeThreshold: 0 })).toThrow(
|
||||
/dedicatedSizeThreshold must be positive/,
|
||||
);
|
||||
expect(() => blob("image", { packFileSizeThreshold: 1.5 })).toThrow(
|
||||
/packFileSizeThreshold must be a safe integer/,
|
||||
);
|
||||
expect(() =>
|
||||
blob("image", { dedicatedSizeThreshold: Number.MAX_SAFE_INTEGER + 1 }),
|
||||
).toThrow(/dedicatedSizeThreshold must be a safe integer/);
|
||||
});
|
||||
});
|
||||
|
||||
class BlobStruct {
|
||||
constructor(
|
||||
public data?: unknown,
|
||||
public uri?: unknown,
|
||||
) {}
|
||||
}
|
||||
|
||||
describe("coerceBlobValue", () => {
|
||||
it.each([
|
||||
["Buffer", Buffer.from("x"), { data: Buffer.from("x"), uri: null }],
|
||||
[
|
||||
"Uint8Array",
|
||||
new Uint8Array([120]),
|
||||
{ data: new Uint8Array([120]), uri: null },
|
||||
],
|
||||
["URI string", "s3://bucket/key", { data: null, uri: "s3://bucket/key" }],
|
||||
[
|
||||
"data struct",
|
||||
{ data: Buffer.from("y") },
|
||||
{ data: Buffer.from("y"), uri: null },
|
||||
],
|
||||
[
|
||||
"uri struct",
|
||||
{ uri: "s3://bucket/key" },
|
||||
{ data: null, uri: "s3://bucket/key" },
|
||||
],
|
||||
[
|
||||
"ArrayBuffer",
|
||||
new Uint8Array([120]).buffer,
|
||||
{ data: new Uint8Array([120]), uri: null },
|
||||
],
|
||||
[
|
||||
"URL",
|
||||
new URL("https://example.com/blob"),
|
||||
{ data: null, uri: "https://example.com/blob" },
|
||||
],
|
||||
[
|
||||
"data struct with ArrayBuffer",
|
||||
{ data: new Uint8Array([121]).buffer },
|
||||
{ data: new Uint8Array([121]), uri: null },
|
||||
],
|
||||
[
|
||||
"uri struct with URL",
|
||||
{ uri: new URL("file:///tmp/a.png") },
|
||||
{ data: null, uri: "file:///tmp/a.png" },
|
||||
],
|
||||
[
|
||||
"class instance with Uint8Array data",
|
||||
new BlobStruct(new Uint8Array([122])),
|
||||
{ data: new Uint8Array([122]), uri: null },
|
||||
],
|
||||
[
|
||||
"class instance with ArrayBuffer data",
|
||||
new BlobStruct(new Uint8Array([122]).buffer),
|
||||
{ data: new Uint8Array([122]), uri: null },
|
||||
],
|
||||
[
|
||||
"class instance with URL uri",
|
||||
new BlobStruct(undefined, new URL("s3://bucket/key")),
|
||||
{ data: null, uri: "s3://bucket/key" },
|
||||
],
|
||||
["null", null, null],
|
||||
])("accepts %s", (_name, input, expected) => {
|
||||
expect(coerceBlobValue(input)).toEqual(expected);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["empty URI", "", /uri cannot be empty/],
|
||||
["object without data or uri", { position: 0 }, /data' or 'uri/],
|
||||
[
|
||||
"Int16Array",
|
||||
new Int16Array([1]),
|
||||
/Blob data must be Buffer, Uint8Array, ArrayBuffer, or Blob/,
|
||||
],
|
||||
[
|
||||
"both data and uri",
|
||||
{ data: Buffer.from("y"), uri: "s3://bucket/key" },
|
||||
/exactly one of 'data' or 'uri'/,
|
||||
],
|
||||
[
|
||||
"neither data nor uri",
|
||||
{ data: null, uri: null },
|
||||
/exactly one of 'data' or 'uri'/,
|
||||
],
|
||||
["Blob", new Blob(["x"]), /must be read asynchronously/],
|
||||
["data struct with Blob", { data: new Blob(["x"]) }, /asynchronously/],
|
||||
[
|
||||
"class instance with Blob data",
|
||||
new BlobStruct(new Blob(["x"])),
|
||||
/asynchronously/,
|
||||
],
|
||||
[
|
||||
"unknown struct field",
|
||||
{ data: Buffer.from("y"), mime: "image/png" },
|
||||
/Blob struct values only support .*got 'mime'/,
|
||||
],
|
||||
[
|
||||
"misspelled struct field",
|
||||
{ uri: "s3://bucket/key", dat: Buffer.from("y") },
|
||||
/got 'dat'/,
|
||||
],
|
||||
[
|
||||
"SharedArrayBuffer",
|
||||
new SharedArrayBuffer(2),
|
||||
/Blob column values must be Buffer, Uint8Array, ArrayBuffer, Blob/,
|
||||
],
|
||||
["Date", new Date(0), /Blob column values must be/],
|
||||
])("rejects %s", (_name, input, message) => {
|
||||
expect(() => coerceBlobValue(input)).toThrow(message);
|
||||
});
|
||||
});
|
||||
|
||||
describe("external blob ranges", () => {
|
||||
const uri = "s3://bucket/key";
|
||||
|
||||
it.each([0, 4, 0n, 4n, 9007199254740993n])(
|
||||
"preserves position %s",
|
||||
(position) => {
|
||||
expect(coerceBlobValue({ uri, position, size: 6 })).toEqual({
|
||||
data: null,
|
||||
uri,
|
||||
position: BigInt(position),
|
||||
size: 6n,
|
||||
});
|
||||
},
|
||||
);
|
||||
|
||||
it.each([
|
||||
[{ uri, position: 0 }, /both position and size/],
|
||||
[{ uri, size: 6 }, /both position and size/],
|
||||
[{ uri, position: null, size: 6 }, /both position and size/],
|
||||
[{ uri, position: 0, size: 0 }, /size must be positive/],
|
||||
[{ uri, position: -1, size: 6 }, /position must be non-negative/],
|
||||
[{ uri, position: 0, size: -1n }, /size must be positive/],
|
||||
[{ uri, position: 0.5, size: 6 }, /safe integer or bigint/],
|
||||
[
|
||||
{ uri, position: Number.MAX_SAFE_INTEGER + 1, size: 6 },
|
||||
/safe integer or bigint/,
|
||||
],
|
||||
[{ uri, position: 0, size: Number.NaN }, /safe integer or bigint/],
|
||||
[{ uri, position: "4", size: 6 }, /safe integer or bigint/],
|
||||
[{ uri, position: false, size: 6 }, /safe integer or bigint/],
|
||||
[{ uri, position: 1n << 64n, size: 6 }, /fit in uint64/],
|
||||
[{ uri, position: 0xffffffffffffffffn, size: 1 }, /overflows uint64/],
|
||||
[{ data: Buffer.from("x"), position: 0, size: 6 }, /require a uri/],
|
||||
])("rejects malformed input %p", (value, message) => {
|
||||
expect(() => coerceBlobValue(value)).toThrow(message);
|
||||
});
|
||||
|
||||
it("preserves ranges while normalizing a URL held by getters", () => {
|
||||
class Range {
|
||||
get uri() {
|
||||
return new URL(uri);
|
||||
}
|
||||
get position() {
|
||||
return 4;
|
||||
}
|
||||
get size() {
|
||||
return 6;
|
||||
}
|
||||
}
|
||||
expect(coerceBlobValue(new Range())).toEqual({
|
||||
data: null,
|
||||
uri,
|
||||
position: 4n,
|
||||
size: 6n,
|
||||
});
|
||||
});
|
||||
|
||||
it("writes ranges and mixed whole-file, inline, and null values to Arrow", () => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
const rows = [
|
||||
null,
|
||||
Buffer.from("x"),
|
||||
uri,
|
||||
{ uri, position: 4, size: 6 },
|
||||
].map((image) => ({ image }));
|
||||
const column = makeArrowTable(rows, { schema }).getChild("image")!;
|
||||
expect(column.get(0)).toBeNull();
|
||||
expect(Buffer.from(column.getChild("data")!.get(1))).toEqual(
|
||||
Buffer.from("x"),
|
||||
);
|
||||
expect(column.getChild("uri")!.get(2)).toBe(uri);
|
||||
expect(column.getChild("position")!.toArray()).toEqual(
|
||||
new BigUint64Array([0n, 0n, 0n, 4n]),
|
||||
);
|
||||
expect(column.getChild("position")!.get(2)).toBeNull();
|
||||
expect(column.getChild("size")!.get(3)).toBe(6n);
|
||||
});
|
||||
|
||||
it.each([
|
||||
Buffer.from("inline"),
|
||||
uri,
|
||||
{ uri, position: null, size: null },
|
||||
null,
|
||||
])("retains the full declared range layout when seeded with %p", (image) => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
const table = makeArrowTable([{ image }], { schema });
|
||||
expect(table.schema).toEqual(schema);
|
||||
expect(table.getChild("image")!.getChild("position")!.get(0)).toBeNull();
|
||||
});
|
||||
|
||||
it.each([0, 4])(
|
||||
"rejects position %s on a minimal schema before dropping keys",
|
||||
(position) => {
|
||||
expect(() =>
|
||||
makeArrowTable(
|
||||
[
|
||||
{ image: Buffer.from("inline") },
|
||||
{ image: { uri, position, size: 6 } },
|
||||
],
|
||||
{ schema: new Schema([legacyBlob("image")]) },
|
||||
),
|
||||
).toThrow(
|
||||
/field image at row 1: Blob storage schema cannot store byte ranges/,
|
||||
);
|
||||
},
|
||||
);
|
||||
|
||||
it("accepts omitted or null range fields on the minimal schema", () => {
|
||||
const column = makeArrowTable(
|
||||
[{ image: { uri, position: null, size: null } }],
|
||||
{
|
||||
schema: new Schema([legacyBlob("image")]),
|
||||
},
|
||||
).getChild("image")!;
|
||||
expect(column.getChild("uri")!.get(0)).toBe(uri);
|
||||
});
|
||||
|
||||
it("validates ranges nested in structs and lists against their storage fields", () => {
|
||||
const value = { uri, position: 4, size: 6 };
|
||||
for (const field of [legacyBlob("image"), blob("image")]) {
|
||||
const schema = new Schema([
|
||||
new Field("info", new Struct([field]), true),
|
||||
new Field("images", new List(field), true),
|
||||
]);
|
||||
if ((field.type as Struct).children.length === 2) {
|
||||
expect(() =>
|
||||
makeArrowTable([{ info: { image: value } }], { schema }),
|
||||
).toThrow(/field info.image.*storage schema/);
|
||||
expect(() => makeArrowTable([{ images: [value] }], { schema })).toThrow(
|
||||
/field images\[0\].*storage schema/,
|
||||
);
|
||||
} else {
|
||||
const table = makeArrowTable(
|
||||
[{ info: { image: value }, images: [value] }],
|
||||
{ schema },
|
||||
);
|
||||
const info = table.getChild("info")!.get(0) as {
|
||||
image: { position: bigint; size: bigint };
|
||||
};
|
||||
expect(info.image.position).toBe(4n);
|
||||
const images = Array.from(
|
||||
table.getChild("images")!.get(0) as Iterable<{
|
||||
position: bigint;
|
||||
size: bigint;
|
||||
}>,
|
||||
);
|
||||
expect(images[0].size).toBe(6n);
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe("resolveBlobInputs", () => {
|
||||
it("reads Blob and File values in blob columns only", async () => {
|
||||
const schema = new Schema([
|
||||
new Field("name", new Utf8(), true),
|
||||
blob("image"),
|
||||
new Field("images", new List(blob("image")), true),
|
||||
]);
|
||||
const shared = new Blob(["same"]);
|
||||
const notABlobColumn = new Blob(["other"]);
|
||||
const data = [
|
||||
{
|
||||
name: notABlobColumn,
|
||||
image: new File(["file"], "a.txt"),
|
||||
images: [shared, { data: shared }, null],
|
||||
},
|
||||
{ name: "x", image: Buffer.from("buf"), images: null },
|
||||
];
|
||||
const resolved = await resolveBlobInputs(data, schema);
|
||||
|
||||
expect(resolved[0].name).toBe(notABlobColumn);
|
||||
expect(Buffer.from(resolved[0].image as Uint8Array).toString()).toBe(
|
||||
"file",
|
||||
);
|
||||
const images = resolved[0].images as [Uint8Array, { data: Uint8Array }];
|
||||
expect(Buffer.from(images[0]).toString()).toBe("same");
|
||||
expect(Buffer.from(images[1].data).toString()).toBe("same");
|
||||
// Rows without Blob values and the caller's records are left alone.
|
||||
expect(resolved[1]).toBe(data[1]);
|
||||
expect(data[0].images?.[0]).toBe(shared);
|
||||
});
|
||||
|
||||
it("reads each distinct Blob once, bare or inside data structs", async () => {
|
||||
class CountingBlob extends Blob {
|
||||
calls = 0;
|
||||
async arrayBuffer(): Promise<ArrayBuffer> {
|
||||
this.calls++;
|
||||
return super.arrayBuffer();
|
||||
}
|
||||
}
|
||||
const shared = new CountingBlob(["same"]);
|
||||
const schema = new Schema([blob("image")]);
|
||||
const resolved = await resolveBlobInputs(
|
||||
[
|
||||
{ image: { data: shared } },
|
||||
{ image: { data: shared } },
|
||||
{ image: shared },
|
||||
],
|
||||
schema,
|
||||
);
|
||||
expect(shared.calls).toBe(1);
|
||||
const first = (resolved[0].image as { data: Uint8Array }).data;
|
||||
expect(Buffer.from(first).toString()).toBe("same");
|
||||
expect((resolved[1].image as { data: Uint8Array }).data).toBe(first);
|
||||
expect(resolved[2].image).toBe(first);
|
||||
});
|
||||
|
||||
it("reads a Blob held by a class instance", async () => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
const [row] = await resolveBlobInputs(
|
||||
[{ image: new BlobStruct(new Blob(["hi"])) }],
|
||||
schema,
|
||||
);
|
||||
expect(
|
||||
Buffer.from((row.image as { data: Uint8Array }).data).toString(),
|
||||
).toBe("hi");
|
||||
});
|
||||
|
||||
it("returns the input when there is nothing to read", async () => {
|
||||
const data = [{ image: Buffer.from("x") }];
|
||||
const schema = new Schema([blob("image")]);
|
||||
await expect(resolveBlobInputs(data, schema)).resolves.toBe(data);
|
||||
await expect(resolveBlobInputs(data)).resolves.toBe(data);
|
||||
});
|
||||
});
|
||||
|
||||
describe("convertToTable blob columns with embedding functions", () => {
|
||||
// Embedding metadata routes rows through ensureNestedFieldsExist, which must
|
||||
// leave blob values for the blob coercion instead of treating them as structs.
|
||||
class LengthEmbedding extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 2;
|
||||
}
|
||||
embeddingDataType(): Float {
|
||||
return new Float32();
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return data.map((text) => [text.length, 1]);
|
||||
}
|
||||
}
|
||||
getRegistry().register("blob-test-length")(LengthEmbedding);
|
||||
const func = getRegistry().get<LengthEmbedding>("blob-test-length")!.create();
|
||||
const embeddingSchema = LanceSchema({
|
||||
text: func.sourceField(new Utf8()),
|
||||
vector: func.vectorField(),
|
||||
});
|
||||
const schema = new Schema(
|
||||
[...embeddingSchema.fields, blob("image")],
|
||||
embeddingSchema.metadata,
|
||||
);
|
||||
|
||||
it.each([
|
||||
["Uint8Array", new Uint8Array([104, 105]), "hi", null],
|
||||
["ArrayBuffer", new Uint8Array([104, 105]).buffer, "hi", null],
|
||||
["Blob", new Blob(["hi"]), "hi", null],
|
||||
["URI string", "s3://b/k", null, "s3://b/k"],
|
||||
["URL", new URL("s3://b/k"), null, "s3://b/k"],
|
||||
["data struct", { data: new Uint8Array([104, 105]) }, "hi", null],
|
||||
])("accepts %s input", async (_name, image, data, uri) => {
|
||||
const table = await convertToTable([{ text: "a", image }], undefined, {
|
||||
schema,
|
||||
});
|
||||
const column = table.getChild("image")!;
|
||||
expect(column.nullCount).toBe(0);
|
||||
const bytes = column.getChild("data")!.get(0);
|
||||
expect(bytes == null ? null : Buffer.from(bytes).toString()).toBe(data);
|
||||
expect(column.getChild("uri")!.get(0)).toBe(uri);
|
||||
});
|
||||
|
||||
it("writes a missing blob as null", async () => {
|
||||
const table = await convertToTable([{ text: "a" }], undefined, { schema });
|
||||
const column = table.getChild("image")!;
|
||||
expect(column.nullCount).toBe(1);
|
||||
expect(column.get(0)).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("BlobFile", () => {
|
||||
it("rejects constructing BlobFile without a native handle", () => {
|
||||
expect(() => new (BlobFile as unknown as { new (): BlobFile })()).toThrow(
|
||||
/fetchBlobFiles/,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("makeArrowTable blob columns", () => {
|
||||
it.each([
|
||||
["ReadableStream", new ReadableStream()],
|
||||
["Int16Array", new Int16Array([1])],
|
||||
])("rejects %s input rather than writing a null blob", (_name, image) => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
expect(() => makeArrowTable([{ image }], { schema })).toThrow(/image/);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["Blob", new Blob(["hello"])],
|
||||
["File", new File(["hello"], "hello.txt")],
|
||||
["data struct with Blob", { data: new Blob(["hello"]) }],
|
||||
[
|
||||
"Blob with enumerable fields",
|
||||
Object.assign(new Blob(["hello"]), { position: 0 }),
|
||||
],
|
||||
])("points %s input at the async APIs", (_name, image) => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
expect(() => makeArrowTable([{ id: 1, image }], { schema })).toThrow(
|
||||
/blob field image at row 0: Blob and File values must be read asynchronously/,
|
||||
);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["ArrayBuffer", new Uint8Array([104, 105]).buffer, "hi", null],
|
||||
[
|
||||
"data struct with ArrayBuffer",
|
||||
{ data: new Uint8Array([104, 105]).buffer },
|
||||
"hi",
|
||||
null,
|
||||
],
|
||||
["URL", new URL("https://example.com/a"), null, "https://example.com/a"],
|
||||
[
|
||||
"uri struct with URL",
|
||||
{ uri: new URL("file:///tmp/a.png") },
|
||||
null,
|
||||
"file:///tmp/a.png",
|
||||
],
|
||||
])("accepts %s input", (_name, image, data, uri) => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
const input = { image };
|
||||
const table = makeArrowTable([input], { schema });
|
||||
const column = table.getChild("image")!;
|
||||
expect(column.nullCount).toBe(0);
|
||||
const bytes = column.getChild("data")!.get(0);
|
||||
expect(bytes == null ? null : Buffer.from(bytes).toString()).toBe(data);
|
||||
expect(column.getChild("uri")!.get(0)).toBe(uri);
|
||||
// The caller's record is left untouched.
|
||||
expect(input.image).toBe(image);
|
||||
});
|
||||
|
||||
it.each([
|
||||
["Uint8Array", new Uint8Array([104, 105]), "hi", null],
|
||||
["URI string", "s3://b/k", null, "s3://b/k"],
|
||||
["data struct", { data: new Uint8Array([104, 105]) }, "hi", null],
|
||||
])("accepts %s input after a null row", (_name, image, data, uri) => {
|
||||
const schema = new Schema([blob("image")]);
|
||||
const table = makeArrowTable([{ image: null }, { image }], { schema });
|
||||
const column = table.getChild("image")!;
|
||||
expect(column.nullCount).toBe(1);
|
||||
expect(column.get(0)).toBeNull();
|
||||
const bytes = column.getChild("data")!.get(1);
|
||||
expect(bytes == null ? null : Buffer.from(bytes).toString()).toBe(data);
|
||||
expect(column.getChild("uri")!.get(1)).toBe(uri);
|
||||
});
|
||||
|
||||
it("accepts a null struct before one holding a blob", () => {
|
||||
const schema = new Schema([
|
||||
new Field("info", new Struct([blob("image")]), true),
|
||||
]);
|
||||
const table = makeArrowTable(
|
||||
[{ info: null }, { info: { image: new Uint8Array([97]) } }],
|
||||
{ schema },
|
||||
);
|
||||
const info = table.getChild("info")!;
|
||||
expect(info.get(0)).toBeNull();
|
||||
const image = (info.get(1) as { image: { data: Uint8Array } }).image;
|
||||
expect(Buffer.from(image.data).toString()).toBe("a");
|
||||
});
|
||||
|
||||
it("skips a missing blob column named like an Object.prototype member", () => {
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int32(), true),
|
||||
blob("constructor"),
|
||||
]);
|
||||
const table = makeArrowTable([{ id: 1 }], { schema });
|
||||
expect(table.schema.fields.map((f) => f.name)).toEqual(["id"]);
|
||||
});
|
||||
|
||||
it("accepts ArrayBuffer and URL inside lists and structs", () => {
|
||||
const schema = new Schema([
|
||||
new Field("images", new List(blob("image")), true),
|
||||
new Field("info", new Struct([blob("image")]), true),
|
||||
]);
|
||||
const table = makeArrowTable(
|
||||
[
|
||||
{
|
||||
images: [new Uint8Array([97]).buffer, new URL("s3://b/k")],
|
||||
info: { image: new Uint8Array([98]).buffer },
|
||||
},
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
const images = Array.from(
|
||||
table.getChild("images")!.get(0) as Iterable<{
|
||||
data: Uint8Array | null;
|
||||
uri: string | null;
|
||||
}>,
|
||||
);
|
||||
expect(Buffer.from(images[0].data!).toString()).toBe("a");
|
||||
expect(images[1].uri).toBe("s3://b/k");
|
||||
const info = table.getChild("info")!.get(0) as {
|
||||
image: { data: Uint8Array };
|
||||
};
|
||||
expect(Buffer.from(info.image.data).toString()).toBe("b");
|
||||
});
|
||||
|
||||
it("coerces Buffer input onto a blob field", () => {
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
blob("image"),
|
||||
]);
|
||||
const table = makeArrowTable([{ id: 1n, image: Buffer.from("hello") }], {
|
||||
schema,
|
||||
});
|
||||
expect(isBlobField(table.schema.fields[1])).toBe(true);
|
||||
const image = table.getChild("image")!;
|
||||
expect(image.nullCount).toBe(0);
|
||||
expect(image.getChild("uri")!.get(0)).toBeNull();
|
||||
expect(image.getChild("data")!.nullCount).toBe(0);
|
||||
expect(Buffer.from(image.getChild("data")!.get(0)!).toString()).toBe(
|
||||
"hello",
|
||||
);
|
||||
});
|
||||
|
||||
it("coerces Buffer elements inside a list and keeps null slots", () => {
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
new Field("images", new List(blob("image")), true),
|
||||
]);
|
||||
const table = makeArrowTable(
|
||||
[
|
||||
{ id: 1n, images: [Buffer.from("a"), Buffer.from("bb")] },
|
||||
{ id: 2n, images: null },
|
||||
{ id: 3n, images: [Buffer.from("c"), null] },
|
||||
{ id: 4n, images: [] },
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
const images = table.getChild("images")!;
|
||||
expect(images.nullCount).toBe(1);
|
||||
const rows = images.toArray();
|
||||
expect(rows[1]).toBeNull();
|
||||
expect(Array.from(rows[3] as Iterable<unknown>)).toHaveLength(0);
|
||||
const first = Array.from(rows[0] as Iterable<{ data: Uint8Array | null }>);
|
||||
expect(Buffer.from(first[0].data!).toString()).toBe("a");
|
||||
expect(Buffer.from(first[1].data!).toString()).toBe("bb");
|
||||
const third = Array.from(
|
||||
rows[2] as Iterable<{ data: Uint8Array | null } | null>,
|
||||
);
|
||||
expect(Buffer.from(third[0]!.data!).toString()).toBe("c");
|
||||
expect(third[1]).toBeNull();
|
||||
});
|
||||
|
||||
it("coerces Buffer fields inside list structs", () => {
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
new Field(
|
||||
"items",
|
||||
new List(
|
||||
new Field(
|
||||
"item",
|
||||
new Struct([new Field("name", new Utf8(), true), blob("image")]),
|
||||
true,
|
||||
),
|
||||
),
|
||||
true,
|
||||
),
|
||||
]);
|
||||
const table = makeArrowTable(
|
||||
[
|
||||
{
|
||||
id: 1n,
|
||||
items: [{ name: "one", image: Buffer.from("alpha") }],
|
||||
},
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
const items = Array.from(
|
||||
table.getChild("items")!.toArray()[0] as Iterable<{
|
||||
name: string;
|
||||
image: { data: Uint8Array | null };
|
||||
}>,
|
||||
);
|
||||
expect(items[0].name).toBe("one");
|
||||
expect(Buffer.from(items[0].image.data!).toString()).toBe("alpha");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,306 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as http from "http";
|
||||
import { Catalog, connectCatalog } from "../lancedb";
|
||||
|
||||
type RecordedRequest = {
|
||||
url: string;
|
||||
headers: http.IncomingHttpHeaders;
|
||||
body: Record<string, unknown>;
|
||||
};
|
||||
|
||||
async function withCatalog(
|
||||
responses: [number, unknown][],
|
||||
callback: (catalog: Catalog, requests: RecordedRequest[]) => Promise<void>,
|
||||
) {
|
||||
const requests: RecordedRequest[] = [];
|
||||
const server = http.createServer(async (req, res) => {
|
||||
const chunks: Buffer[] = [];
|
||||
for await (const chunk of req) chunks.push(Buffer.from(chunk));
|
||||
const body = Buffer.concat(chunks).toString();
|
||||
requests.push({
|
||||
url: req.url ?? "",
|
||||
headers: req.headers,
|
||||
body: body ? JSON.parse(body) : {},
|
||||
});
|
||||
const [status, response] = responses.shift() ?? [
|
||||
500,
|
||||
{ error: "Unexpected request" },
|
||||
];
|
||||
res.writeHead(status, { "content-type": "application/json" });
|
||||
res.end(status === 204 ? undefined : JSON.stringify(response));
|
||||
});
|
||||
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
|
||||
const address = server.address();
|
||||
if (!address || typeof address === "string")
|
||||
throw new Error("Missing server address");
|
||||
try {
|
||||
const catalog = await connectCatalog(`http://127.0.0.1:${address.port}`, {
|
||||
apiKey: "secret",
|
||||
sqlHostOverride: "grpc+tls://sql.example.com:10026",
|
||||
clientConfig: {
|
||||
extraHeaders: {
|
||||
"X-LanceDB-Database": "wrong-static",
|
||||
"X-LanceDB-Database-Prefix": "wrong",
|
||||
},
|
||||
},
|
||||
headerProvider: () => ({
|
||||
"X-LanceDB-Database": "wrong-dynamic",
|
||||
"X-LanceDB-Database-Prefix": "wrong",
|
||||
authorization: "Bearer refreshed",
|
||||
}),
|
||||
});
|
||||
await callback(catalog, requests);
|
||||
expect(responses).toHaveLength(0);
|
||||
} finally {
|
||||
server.closeAllConnections();
|
||||
await new Promise<void>((resolve) => server.close(() => resolve()));
|
||||
}
|
||||
}
|
||||
|
||||
describe("remote catalog", () => {
|
||||
it("uses root namespace routes and preserves independent database scope", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
[204, null],
|
||||
[200, { tables: [] }],
|
||||
[200, {}],
|
||||
[200, { tables: [] }],
|
||||
[200, { tables: [] }],
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["team/search"], page_token: "next" }],
|
||||
[204, null],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const first = await catalog.createDatabase("team/search", {
|
||||
existOk: true,
|
||||
});
|
||||
expect(await first.tableNames()).toEqual([]);
|
||||
const second = await catalog.connectDatabase("other");
|
||||
expect(await second.tableNames()).toEqual([]);
|
||||
expect(await first.tableNames()).toEqual([]);
|
||||
const names = catalog.listDatabases({ pageLimit: 1, pageToken: "a/b" });
|
||||
expect(await names.next()).toEqual({
|
||||
done: false,
|
||||
value: "team/search",
|
||||
});
|
||||
expect(names.pageToken()).toBe("next");
|
||||
await catalog.dropDatabase("team/search", { ignoreMissing: true });
|
||||
expect(requests[0].url).toBe("/v1/namespace/team%2Fsearch/create");
|
||||
expect(requests[0].body).toEqual({ mode: "ExistOk" });
|
||||
expect(requests[5].url).toBe(
|
||||
"/v1/namespace/$/list?limit=1&page_token=a%2Fb",
|
||||
);
|
||||
expect(requests[6].body).toEqual({
|
||||
mode: "Skip",
|
||||
behavior: "Restrict",
|
||||
});
|
||||
for (const [i, request] of requests.entries()) {
|
||||
expect(request.headers["x-lancedb-database"]).toBe(
|
||||
i === 1 || i === 4 ? "team/search" : i === 3 ? "other" : undefined,
|
||||
);
|
||||
expect(request.headers["x-lancedb-database-prefix"]).toBeUndefined();
|
||||
expect(request.headers.authorization).toBe("Bearer refreshed");
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("propagates lifecycle errors and sends restricted drops", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
[404, {}],
|
||||
[409, {}],
|
||||
[400, {}],
|
||||
[404, {}],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
await expect(catalog.connectDatabase("missing")).rejects.toThrow(
|
||||
"missing",
|
||||
);
|
||||
await expect(catalog.createDatabase("exists")).rejects.toThrow(
|
||||
"exists",
|
||||
);
|
||||
await expect(catalog.dropDatabase("full")).rejects.toThrow();
|
||||
await catalog.dropDatabase("missing", { ignoreMissing: true });
|
||||
expect(requests[2].body).toEqual({
|
||||
mode: "Fail",
|
||||
behavior: "Restrict",
|
||||
});
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("lazily traverses pages, including empty pages, and exposes cached state", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["a", "b"], page_token: "empty" }],
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: [], page_token: "last" }],
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["c", "d"], page_token: "" }],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const names = catalog.listDatabases({ pageLimit: 2 });
|
||||
expect(names[Symbol.asyncIterator]()).toBe(names);
|
||||
expect(requests).toHaveLength(0);
|
||||
expect(names.numPageResults()).toBe(0);
|
||||
expect(names.pageToken()).toBeUndefined();
|
||||
expect(await names.next()).toEqual({ done: false, value: "a" });
|
||||
expect(names.numPageResults()).toBe(1);
|
||||
expect(names.pageToken()).toBe("empty");
|
||||
expect(await names.next()).toEqual({ done: false, value: "b" });
|
||||
expect(requests).toHaveLength(1);
|
||||
expect(await names.next()).toEqual({ done: false, value: "c" });
|
||||
expect(requests).toHaveLength(3);
|
||||
expect(names.pageToken()).toBeUndefined();
|
||||
expect(names.numPageResults()).toBe(1);
|
||||
const remaining = [];
|
||||
for await (const name of names) remaining.push(name);
|
||||
expect(remaining).toEqual(["d"]);
|
||||
expect(names.numPageResults()).toBe(0);
|
||||
expect(await names.next()).toEqual({ done: true, value: undefined });
|
||||
expect(requests.map((request) => request.url)).toEqual([
|
||||
"/v1/namespace/$/list?limit=2",
|
||||
"/v1/namespace/$/list?limit=2&page_token=empty",
|
||||
"/v1/namespace/$/list?limit=2&page_token=last",
|
||||
]);
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("resumes from a saved token after draining the cache", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["a", "b"], page_token: "next/token" }],
|
||||
[200, { namespaces: ["c"] }],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const names = catalog.listDatabases({
|
||||
pageToken: "start/token",
|
||||
pageLimit: 2,
|
||||
});
|
||||
expect(names.pageToken()).toBe("start/token");
|
||||
expect(await names.next()).toEqual({ done: false, value: "a" });
|
||||
while (names.numPageResults() > 0) await names.next();
|
||||
const resumed = catalog.listDatabases({
|
||||
pageToken: names.pageToken(),
|
||||
pageLimit: 2,
|
||||
});
|
||||
const results = [];
|
||||
for await (const name of resumed) results.push(name);
|
||||
expect(results).toEqual(["c"]);
|
||||
expect(requests.map((request) => request.url)).toEqual([
|
||||
"/v1/namespace/$/list?limit=2&page_token=start%2Ftoken",
|
||||
"/v1/namespace/$/list?limit=2&page_token=next%2Ftoken",
|
||||
]);
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("terminates on errors and retains the failed request token", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
[200, { namespaces: ["a"], page_token: "retry" }],
|
||||
[400, {}],
|
||||
[200, { namespaces: ["b"] }],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const names = catalog.listDatabases();
|
||||
await names.next();
|
||||
await expect(names.next()).rejects.toThrow();
|
||||
expect(names.pageToken()).toBe("retry");
|
||||
expect(names.numPageResults()).toBe(0);
|
||||
expect(await names.next()).toEqual({ done: true, value: undefined });
|
||||
expect(requests).toHaveLength(2);
|
||||
const resumed = catalog.listDatabases({ pageToken: names.pageToken() });
|
||||
expect(await resumed.next()).toEqual({ done: false, value: "b" });
|
||||
expect(await resumed.next()).toEqual({ done: true, value: undefined });
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("handles an empty listing and concurrent advances", async () => {
|
||||
await withCatalog(
|
||||
[
|
||||
[200, { namespaces: [] }],
|
||||
[200, { namespaces: ["a", "b"] }],
|
||||
],
|
||||
async (catalog, requests) => {
|
||||
const empty = catalog.listDatabases();
|
||||
expect(await empty.next()).toEqual({ done: true, value: undefined });
|
||||
expect(await empty.next()).toEqual({ done: true, value: undefined });
|
||||
const names = catalog.listDatabases();
|
||||
const pending = [names.next(), names.next(), names.next()];
|
||||
expect(names.numPageResults()).toBe(0);
|
||||
expect(names.pageToken()).toBeUndefined();
|
||||
const results = await Promise.all(pending);
|
||||
expect(
|
||||
results
|
||||
.filter((result) => !result.done)
|
||||
.map((result) => result.value)
|
||||
.sort(),
|
||||
).toEqual(["a", "b"]);
|
||||
expect(results.filter((result) => result.done)).toHaveLength(1);
|
||||
expect(requests).toHaveLength(2);
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
it("validates endpoints and pagination", async () => {
|
||||
await expect(connectCatalog("/tmp/catalog")).rejects.toThrow();
|
||||
const catalog = await connectCatalog("http://127.0.0.1:1");
|
||||
for (const limit of [
|
||||
0,
|
||||
-1,
|
||||
1.5,
|
||||
2147483648,
|
||||
Number.NaN,
|
||||
Number.POSITIVE_INFINITY,
|
||||
]) {
|
||||
expect(() => catalog.listDatabases({ pageLimit: limit })).toThrow(
|
||||
"limit",
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
it.each([
|
||||
"",
|
||||
".",
|
||||
"..",
|
||||
"/db",
|
||||
"db/",
|
||||
"a//b",
|
||||
"a/./b",
|
||||
"a/../b",
|
||||
"a b",
|
||||
"a:b",
|
||||
"a$b",
|
||||
"a%b",
|
||||
"a?b",
|
||||
"a#b",
|
||||
"a\\b",
|
||||
"café",
|
||||
"a\nb",
|
||||
])(
|
||||
"rejects invalid database name %j before sending requests",
|
||||
async (name) => {
|
||||
await withCatalog([], async (catalog, requests) => {
|
||||
for (const method of [
|
||||
"createDatabase",
|
||||
"connectDatabase",
|
||||
"dropDatabase",
|
||||
] as const) {
|
||||
await expect(catalog[method](name)).rejects.toThrow(
|
||||
"Invalid database name",
|
||||
);
|
||||
}
|
||||
expect(requests).toEqual([]);
|
||||
});
|
||||
},
|
||||
);
|
||||
});
|
||||
@@ -1,10 +1,17 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { readdirSync } from "fs";
|
||||
import { mkdirSync, readdirSync } from "fs";
|
||||
import { join } from "path";
|
||||
import { Field, Float64, Schema } from "apache-arrow";
|
||||
import * as tmp from "tmp";
|
||||
import { Connection, Table, connect } from "../lancedb";
|
||||
import {
|
||||
Connection,
|
||||
ListTablesResponse,
|
||||
Table,
|
||||
connect,
|
||||
connectNamespace,
|
||||
} from "../lancedb";
|
||||
import { LocalTable } from "../lancedb/table";
|
||||
|
||||
describe("when connecting", () => {
|
||||
@@ -47,6 +54,15 @@ describe("given a connection", () => {
|
||||
await db.close();
|
||||
expect(db.isOpen()).toBe(false);
|
||||
await expect(db.tableNames()).rejects.toThrow("Connection is closed");
|
||||
await expect(db.listTables()).rejects.toThrow("Connection is closed");
|
||||
await expect(db.renameTable("a", "b")).rejects.toThrow(
|
||||
"Connection is closed",
|
||||
);
|
||||
});
|
||||
|
||||
it("should report renameTable as unsupported on an OSS connection", async () => {
|
||||
await db.createTable("a", [{ id: 1 }]);
|
||||
await expect(db.renameTable("a", "b")).rejects.toThrow(/not supported/);
|
||||
});
|
||||
it("should be able to create a table from an object arg `createTable(options)`, or args `createTable(name, data, options)`", async () => {
|
||||
let tbl = await db.createTable("test", [{ id: 1 }, { id: 2 }]);
|
||||
@@ -81,6 +97,16 @@ describe("given a connection", () => {
|
||||
await db.createTable("test4", [{ id: 1 }, { id: 2 }]);
|
||||
});
|
||||
|
||||
it("should return a completed job when dropping a local table", async () => {
|
||||
await db.createTable("async-drop", [{ id: 1 }]);
|
||||
|
||||
const job = await db.dropTableAsync("async-drop");
|
||||
expect(job.id).toBeNull();
|
||||
await expect(job.status()).resolves.toBe("finished");
|
||||
await job.wait();
|
||||
await expect(db.tableNames()).resolves.toEqual([]);
|
||||
});
|
||||
|
||||
it("should fail if creating table twice, unless overwrite is true", async () => {
|
||||
let tbl = await db.createTable("test", [{ id: 1 }, { id: 2 }]);
|
||||
await expect(tbl.countRows()).resolves.toBe(2);
|
||||
@@ -111,6 +137,92 @@ describe("given a connection", () => {
|
||||
expect(tables).toEqual(["b", "c"]);
|
||||
});
|
||||
|
||||
it("should respect limit and page token when listing tables", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
|
||||
await db.createTable("b", [{ id: 1 }]);
|
||||
await db.createTable("a", [{ id: 1 }]);
|
||||
await db.createTable("c", [{ id: 1 }]);
|
||||
|
||||
const all = await db.listTables();
|
||||
expect(all.tables).toEqual(["a", "b", "c"]);
|
||||
expect(all.pageToken).toBeUndefined();
|
||||
|
||||
const first = await db.listTables({ limit: 1 });
|
||||
expect(first.tables).toEqual(["a"]);
|
||||
expect(first.pageToken).toBeDefined();
|
||||
|
||||
const second = await db.listTables({
|
||||
limit: 1,
|
||||
pageToken: first.pageToken,
|
||||
});
|
||||
expect(second.tables).toEqual(["b"]);
|
||||
});
|
||||
|
||||
it("should visit every table exactly once when walking pages", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
|
||||
const created = ["a", "b", "c", "d", "e"];
|
||||
for (const name of created) {
|
||||
await db.createTable(name, [{ id: 1 }]);
|
||||
}
|
||||
|
||||
const seen: string[] = [];
|
||||
let pageToken: string | undefined = undefined;
|
||||
do {
|
||||
const page: ListTablesResponse = await db.listTables({
|
||||
limit: 2,
|
||||
pageToken,
|
||||
});
|
||||
seen.push(...page.tables);
|
||||
pageToken = page.pageToken;
|
||||
} while (pageToken);
|
||||
|
||||
expect(seen).toEqual(created);
|
||||
});
|
||||
|
||||
it("should list all names and use a default page size of 100", async () => {
|
||||
const names = Array.from(
|
||||
{ length: 130 },
|
||||
(_, i) => `t${i.toString().padStart(3, "0")}`,
|
||||
);
|
||||
for (const name of names) {
|
||||
mkdirSync(join(tmpDir.name, `${name}.lance`));
|
||||
}
|
||||
await expect(db.tableNames()).resolves.toEqual(names);
|
||||
await expect(db.tableNames({ limit: 0 })).resolves.toEqual([]);
|
||||
const first = await db.listTables();
|
||||
expect(first.tables).toEqual(names.slice(0, 100));
|
||||
expect(first.pageToken).toBeDefined();
|
||||
const second = await db.listTables({ pageToken: first.pageToken });
|
||||
expect(second.tables).toEqual(names.slice(100));
|
||||
expect(second.pageToken).toBeUndefined();
|
||||
await expect(db.listTables({ limit: 200 })).resolves.toEqual({
|
||||
tables: names,
|
||||
pageToken: undefined,
|
||||
});
|
||||
await expect(db.listTables({ limit: 0 })).resolves.toEqual({
|
||||
tables: [],
|
||||
pageToken: undefined,
|
||||
});
|
||||
});
|
||||
|
||||
it("should list tables in a namespace", async () => {
|
||||
const db = await connect(tmpDir.name, {
|
||||
// biome-ignore lint/style/useNamingConvention: opaque backend property key, must match Rust
|
||||
namespaceClientProperties: { manifest_enabled: "true" },
|
||||
});
|
||||
await db.createNamespace(["child"]);
|
||||
await db.createTable("nested", [{ id: 1 }], ["child"]);
|
||||
|
||||
await expect(db.listTables(["child"])).resolves.toEqual(
|
||||
expect.objectContaining({ tables: ["nested"] }),
|
||||
);
|
||||
await expect(db.listTables()).resolves.toEqual(
|
||||
expect.objectContaining({ tables: [] }),
|
||||
);
|
||||
});
|
||||
|
||||
it("should create tables in v2 mode", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [...Array(10000).keys()].map((i) => ({ id: i }));
|
||||
@@ -163,18 +275,22 @@ describe("given a connection", () => {
|
||||
|
||||
let manifestDir =
|
||||
tmpDir.name + "/test_manifest_paths_v2_empty.lance/_versions";
|
||||
readdirSync(manifestDir).forEach((file) => {
|
||||
expect(file).toMatch(/^\d{20}\.manifest$/);
|
||||
});
|
||||
readdirSync(manifestDir)
|
||||
.filter((f) => f.endsWith(".manifest"))
|
||||
.forEach((file) => {
|
||||
expect(file).toMatch(/^\d{20}\.manifest$/);
|
||||
});
|
||||
|
||||
table = (await db.createTable("test_manifest_paths_v2", [{ id: 1 }], {
|
||||
enableV2ManifestPaths: true,
|
||||
})) as LocalTable;
|
||||
expect(await table.usesV2ManifestPaths()).toBe(true);
|
||||
manifestDir = tmpDir.name + "/test_manifest_paths_v2.lance/_versions";
|
||||
readdirSync(manifestDir).forEach((file) => {
|
||||
expect(file).toMatch(/^\d{20}\.manifest$/);
|
||||
});
|
||||
readdirSync(manifestDir)
|
||||
.filter((f) => f.endsWith(".manifest"))
|
||||
.forEach((file) => {
|
||||
expect(file).toMatch(/^\d{20}\.manifest$/);
|
||||
});
|
||||
});
|
||||
|
||||
it("should be able to migrate tables to the V2 manifest paths", async () => {
|
||||
@@ -191,15 +307,305 @@ describe("given a connection", () => {
|
||||
|
||||
const manifestDir =
|
||||
tmpDir.name + "/test_manifest_path_migration.lance/_versions";
|
||||
readdirSync(manifestDir).forEach((file) => {
|
||||
expect(file).toMatch(/^\d\.manifest$/);
|
||||
});
|
||||
readdirSync(manifestDir)
|
||||
.filter((f) => f.endsWith(".manifest"))
|
||||
.forEach((file) => {
|
||||
expect(file).toMatch(/^\d\.manifest$/);
|
||||
});
|
||||
|
||||
await table.migrateManifestPathsV2();
|
||||
expect(await table.usesV2ManifestPaths()).toBe(true);
|
||||
|
||||
readdirSync(manifestDir).forEach((file) => {
|
||||
expect(file).toMatch(/^\d{20}\.manifest$/);
|
||||
readdirSync(manifestDir)
|
||||
.filter((f) => f.endsWith(".manifest"))
|
||||
.forEach((file) => {
|
||||
expect(file).toMatch(/^\d{20}\.manifest$/);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("clone table functionality", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let db: Connection;
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
db = await connect(tmpDir.name);
|
||||
});
|
||||
afterEach(() => tmpDir.removeCallback());
|
||||
|
||||
it("should clone a table with latest version (default behavior)", async () => {
|
||||
// Create source table with some data
|
||||
const data = [
|
||||
{ id: 1, text: "hello", vector: [1.0, 2.0] },
|
||||
{ id: 2, text: "world", vector: [3.0, 4.0] },
|
||||
];
|
||||
const sourceTable = await db.createTable("source", data);
|
||||
|
||||
// Add more data to create a new version
|
||||
const moreData = [{ id: 3, text: "test", vector: [5.0, 6.0] }];
|
||||
await sourceTable.add(moreData);
|
||||
|
||||
// Clone the table (should get latest version with 3 rows)
|
||||
const sourceUri = `${tmpDir.name}/source.lance`;
|
||||
const clonedTable = await db.cloneTable("cloned", sourceUri);
|
||||
|
||||
// Verify cloned table has all 3 rows
|
||||
expect(await clonedTable.countRows()).toBe(3);
|
||||
expect((await db.tableNames()).includes("cloned")).toBe(true);
|
||||
});
|
||||
|
||||
it("should clone a table from a specific version", async () => {
|
||||
// Create source table with initial data
|
||||
const data = [
|
||||
{ id: 1, text: "hello", vector: [1.0, 2.0] },
|
||||
{ id: 2, text: "world", vector: [3.0, 4.0] },
|
||||
];
|
||||
const sourceTable = await db.createTable("source", data);
|
||||
|
||||
// Get the initial version
|
||||
const initialVersion = await sourceTable.version();
|
||||
|
||||
// Add more data to create a new version
|
||||
const moreData = [{ id: 3, text: "test", vector: [5.0, 6.0] }];
|
||||
await sourceTable.add(moreData);
|
||||
|
||||
// Verify source now has 3 rows
|
||||
expect(await sourceTable.countRows()).toBe(3);
|
||||
|
||||
// Clone from the initial version (should have only 2 rows)
|
||||
const sourceUri = `${tmpDir.name}/source.lance`;
|
||||
const clonedTable = await db.cloneTable("cloned", sourceUri, {
|
||||
sourceVersion: initialVersion,
|
||||
});
|
||||
|
||||
// Verify cloned table has only the initial 2 rows
|
||||
expect(await clonedTable.countRows()).toBe(2);
|
||||
});
|
||||
|
||||
it("should clone a table from a tagged version", async () => {
|
||||
// Create source table with initial data
|
||||
const data = [
|
||||
{ id: 1, text: "hello", vector: [1.0, 2.0] },
|
||||
{ id: 2, text: "world", vector: [3.0, 4.0] },
|
||||
];
|
||||
const sourceTable = await db.createTable("source", data);
|
||||
|
||||
// Create a tag for the current version
|
||||
const tags = await sourceTable.tags();
|
||||
await tags.create("v1.0", await sourceTable.version());
|
||||
|
||||
// Add more data after the tag
|
||||
const moreData = [{ id: 3, text: "test", vector: [5.0, 6.0] }];
|
||||
await sourceTable.add(moreData);
|
||||
|
||||
// Verify source now has 3 rows
|
||||
expect(await sourceTable.countRows()).toBe(3);
|
||||
|
||||
// Clone from the tagged version (should have only 2 rows)
|
||||
const sourceUri = `${tmpDir.name}/source.lance`;
|
||||
const clonedTable = await db.cloneTable("cloned", sourceUri, {
|
||||
sourceTag: "v1.0",
|
||||
});
|
||||
|
||||
// Verify cloned table has only the tagged version's 2 rows
|
||||
expect(await clonedTable.countRows()).toBe(2);
|
||||
});
|
||||
|
||||
it("should fail when attempting deep clone", async () => {
|
||||
// Create source table with some data
|
||||
const data = [
|
||||
{ id: 1, text: "hello", vector: [1.0, 2.0] },
|
||||
{ id: 2, text: "world", vector: [3.0, 4.0] },
|
||||
];
|
||||
await db.createTable("source", data);
|
||||
|
||||
// Try to create a deep clone (should fail)
|
||||
const sourceUri = `${tmpDir.name}/source.lance`;
|
||||
await expect(
|
||||
db.cloneTable("cloned", sourceUri, { isShallow: false }),
|
||||
).rejects.toThrow("Deep clone is not yet implemented");
|
||||
});
|
||||
});
|
||||
|
||||
describe("namespaces", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let db: Connection;
|
||||
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
// The local DirectoryNamespace backend only supports child namespaces
|
||||
// when manifest mode is enabled (see lance-namespace-impls/src/dir.rs).
|
||||
db = await connect(tmpDir.name, {
|
||||
// biome-ignore lint/style/useNamingConvention: opaque backend property key, must match Rust
|
||||
namespaceClientProperties: { manifest_enabled: "true" },
|
||||
});
|
||||
});
|
||||
afterEach(() => tmpDir.removeCallback());
|
||||
|
||||
it("should create and describe a namespace", async () => {
|
||||
await db.createNamespace(["myns"]);
|
||||
const desc = await db.describeNamespace(["myns"]);
|
||||
expect(desc).toBeDefined();
|
||||
});
|
||||
|
||||
it("should list namespaces created at the root", async () => {
|
||||
await db.createNamespace(["alpha"]);
|
||||
await db.createNamespace(["beta"]);
|
||||
const list = await db.listNamespaces();
|
||||
expect(list.namespaces).toEqual(expect.arrayContaining(["alpha", "beta"]));
|
||||
});
|
||||
|
||||
it("should list child namespaces under a parent", async () => {
|
||||
await db.createNamespace(["parent"]);
|
||||
await db.createNamespace(["parent", "child"]);
|
||||
const list = await db.listNamespaces(["parent"]);
|
||||
expect(list.namespaces).toContain("child");
|
||||
});
|
||||
|
||||
it("should drop a namespace", async () => {
|
||||
await db.createNamespace(["ephemeral"]);
|
||||
await db.dropNamespace(["ephemeral"]);
|
||||
const list = await db.listNamespaces();
|
||||
expect(list.namespaces).not.toContain("ephemeral");
|
||||
});
|
||||
|
||||
it("should raise an error on any namespace op after close", async () => {
|
||||
await db.close();
|
||||
await expect(db.describeNamespace(["foo"])).rejects.toThrow(
|
||||
"Connection is closed",
|
||||
);
|
||||
await expect(db.listNamespaces()).rejects.toThrow("Connection is closed");
|
||||
await expect(db.createNamespace(["foo"])).rejects.toThrow(
|
||||
"Connection is closed",
|
||||
);
|
||||
await expect(db.dropNamespace(["foo"])).rejects.toThrow(
|
||||
"Connection is closed",
|
||||
);
|
||||
});
|
||||
|
||||
it("should raise an understandable error when describing a non-existent namespace", async () => {
|
||||
await expect(db.describeNamespace(["does-not-exist"])).rejects.toThrow(
|
||||
/not found/i,
|
||||
);
|
||||
});
|
||||
|
||||
it("should raise an error when creating a namespace that already exists", async () => {
|
||||
await db.createNamespace(["dup"]);
|
||||
await expect(db.createNamespace(["dup"])).rejects.toThrow();
|
||||
});
|
||||
|
||||
it("should reject an unrecognized createNamespace mode with a clear error", async () => {
|
||||
await expect(
|
||||
// biome-ignore lint/suspicious/noExplicitAny: deliberately bypass TS to test runtime validation
|
||||
db.createNamespace(["x"], { mode: "frobnicate" as any }),
|
||||
).rejects.toThrow(/Invalid mode 'frobnicate'/);
|
||||
});
|
||||
|
||||
it("should reject an unrecognized dropNamespace mode with a clear error", async () => {
|
||||
await db.createNamespace(["x"]);
|
||||
await expect(
|
||||
// biome-ignore lint/suspicious/noExplicitAny: deliberately bypass TS to test runtime validation
|
||||
db.dropNamespace(["x"], { mode: "frobnicate" as any }),
|
||||
).rejects.toThrow(/Invalid mode 'frobnicate'/);
|
||||
});
|
||||
|
||||
it("should reject an unrecognized dropNamespace behavior with a clear error", async () => {
|
||||
await db.createNamespace(["x"]);
|
||||
await expect(
|
||||
// biome-ignore lint/suspicious/noExplicitAny: deliberately bypass TS to test runtime validation
|
||||
db.dropNamespace(["x"], { behavior: "frobnicate" as any }),
|
||||
).rejects.toThrow(/Invalid behavior 'frobnicate'/);
|
||||
});
|
||||
});
|
||||
|
||||
describe("connectNamespace", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
beforeEach(() => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
});
|
||||
afterEach(() => tmpDir.removeCallback());
|
||||
|
||||
it("connects via the dir implementation and supports table ops", async () => {
|
||||
const db = await connectNamespace("dir", { root: tmpDir.name });
|
||||
await db.createTable("users", [{ id: 1 }, { id: 2 }]);
|
||||
await expect(db.tableNames()).resolves.toContain("users");
|
||||
});
|
||||
|
||||
it("throws a clear error when implName is empty", async () => {
|
||||
await expect(connectNamespace("", {})).rejects.toThrow(
|
||||
"implName must be a non-empty string",
|
||||
);
|
||||
});
|
||||
|
||||
it("throws when the namespace implementation is unknown", async () => {
|
||||
await expect(connectNamespace("not-a-real-impl", {})).rejects.toThrow();
|
||||
});
|
||||
|
||||
it("passes storage options through to the namespace", async () => {
|
||||
const db = await connectNamespace(
|
||||
"dir",
|
||||
{ root: tmpDir.name },
|
||||
{ storageOptions: { newTableDataStorageVersion: "stable" } },
|
||||
);
|
||||
await db.createTable("plumbing", [{ id: 1 }]);
|
||||
await expect(db.tableNames()).resolves.toContain("plumbing");
|
||||
});
|
||||
|
||||
it("supports child namespaces when manifestEnabled is true on the dir config", async () => {
|
||||
const writer = await connectNamespace("dir", {
|
||||
root: tmpDir.name,
|
||||
manifestEnabled: true,
|
||||
});
|
||||
await writer.createNamespace(["analytics"]);
|
||||
await writer.createTable("orders", [{ id: 1 }, { id: 2 }], ["analytics"]);
|
||||
await writer.close();
|
||||
|
||||
const reader = await connectNamespace("dir", {
|
||||
root: tmpDir.name,
|
||||
manifestEnabled: true,
|
||||
});
|
||||
await expect(reader.tableNames(["analytics"])).resolves.toContain("orders");
|
||||
const orders = await reader.openTable("orders", ["analytics"]);
|
||||
await expect(orders.countRows()).resolves.toBe(2);
|
||||
});
|
||||
|
||||
it("merges extraProperties into the dir config and is overridden by typed fields", async () => {
|
||||
// Two observable assertions:
|
||||
// - Typed `root` overrides extraProperties.root: createTable would fail
|
||||
// under the bogus path if the override didn't happen.
|
||||
// - extraProperties.manifest_enabled="false" is honored end-to-end. Child
|
||||
// namespaces require manifest mode (default true), so explicitly
|
||||
// disabling it via extraProperties must make createNamespace reject. If
|
||||
// extraProperties pass-through were silently broken, the default would
|
||||
// let createNamespace succeed.
|
||||
const db = await connectNamespace("dir", {
|
||||
root: tmpDir.name,
|
||||
extraProperties: {
|
||||
root: "/should/be/overridden",
|
||||
// biome-ignore lint/style/useNamingConvention: backend property key
|
||||
manifest_enabled: "false",
|
||||
},
|
||||
});
|
||||
await db.createTable("base", [{ id: 1 }]);
|
||||
await expect(db.tableNames()).resolves.toContain("base");
|
||||
await expect(db.createNamespace(["analytics"])).rejects.toThrow();
|
||||
});
|
||||
|
||||
it("flows unknown top-level keys through when implName is dynamic (no silent drop)", async () => {
|
||||
// Routes via the third overload because `impl` is `string`, not the
|
||||
// literal `"dir"`. The dispatcher still notices the runtime value is
|
||||
// "dir", but unknown keys like `manifest_enabled` must not be silently
|
||||
// dropped during the conversion.
|
||||
//
|
||||
// Asserting a *negative* outcome (manifest disabled -> createNamespace
|
||||
// rejects) is required for observability, since the backend default for
|
||||
// `manifest_enabled` is true.
|
||||
const impl: string = "dir";
|
||||
const db = await connectNamespace(impl, {
|
||||
root: tmpDir.name,
|
||||
// biome-ignore lint/style/useNamingConvention: backend property key
|
||||
manifest_enabled: "false",
|
||||
});
|
||||
await expect(db.createNamespace(["mixed"])).rejects.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -11,8 +11,11 @@ import {
|
||||
Float16,
|
||||
Float32,
|
||||
Float64,
|
||||
Int32,
|
||||
Schema,
|
||||
Utf8,
|
||||
fromDataToBuffer,
|
||||
tableFromIPC,
|
||||
} from "../lancedb/arrow";
|
||||
import { EmbeddingFunction, LanceSchema } from "../lancedb/embedding";
|
||||
import { getRegistry, register } from "../lancedb/embedding/registry";
|
||||
@@ -184,6 +187,115 @@ describe("embedding functions", () => {
|
||||
const vector0 = JSON.parse(JSON.stringify(arr[0].vector));
|
||||
expect(vector0).toEqual([1, 2, 3]);
|
||||
});
|
||||
it("should append multiple Python embeddings with the same alias", async () => {
|
||||
@register("python-mock")
|
||||
// biome-ignore lint/correctness/noUnusedVariables: the decorator registers this class
|
||||
class MockEmbeddingFunction extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 3;
|
||||
}
|
||||
embeddingDataType(): Float {
|
||||
return new Float32();
|
||||
}
|
||||
async computeQueryEmbeddings(_data: string) {
|
||||
return [1, 2, 3];
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return data.map((value) =>
|
||||
value === "hello world" ? [1, 2, 3] : [4, 5, 6],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const metadata = new Map([
|
||||
[
|
||||
"embedding_functions",
|
||||
'[{"source_column":"text1","vector_column":"vector1","name":"python-mock","model":{}},{"source_column":"text2","vector_column":"vector2","name":"python-mock","model":{}}]',
|
||||
],
|
||||
]);
|
||||
const schema = new Schema(
|
||||
[
|
||||
new Field("text1", new Utf8(), true),
|
||||
new Field("text2", new Utf8(), true),
|
||||
new Field(
|
||||
"vector1",
|
||||
new FixedSizeList(3, new Field("item", new Float32(), true)),
|
||||
true,
|
||||
),
|
||||
new Field(
|
||||
"vector2",
|
||||
new FixedSizeList(3, new Field("item", new Float32(), true)),
|
||||
true,
|
||||
),
|
||||
],
|
||||
metadata,
|
||||
);
|
||||
|
||||
const db = await connect(tmpDir.name);
|
||||
const table = await db.createEmptyTable("test", schema);
|
||||
await table.add([{ text1: "hello world", text2: "goodbye world" }]);
|
||||
|
||||
const rows = await table.query().toArray();
|
||||
expect(JSON.parse(JSON.stringify(rows[0].vector1))).toEqual([1, 2, 3]);
|
||||
expect(JSON.parse(JSON.stringify(rows[0].vector2))).toEqual([4, 5, 6]);
|
||||
});
|
||||
|
||||
it("should append generated vectors to a non-nullable schema", async () => {
|
||||
@register("non_nullable_schema_test")
|
||||
class MockEmbeddingFunction extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 3;
|
||||
}
|
||||
embeddingDataType(): Float {
|
||||
return new Float64();
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return data.map(() => [1, 2, 3]);
|
||||
}
|
||||
}
|
||||
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int32()),
|
||||
new Field("text", new Utf8()),
|
||||
new Field("type", new Utf8()),
|
||||
new Field(
|
||||
"vector",
|
||||
new FixedSizeList(3, new Field("item", new Float64())),
|
||||
),
|
||||
]);
|
||||
const func = new MockEmbeddingFunction();
|
||||
const db = await connect(tmpDir.name);
|
||||
const table = await db.createEmptyTable("test_non_nullable", schema, {
|
||||
embeddingFunction: {
|
||||
function: func,
|
||||
sourceColumn: "text",
|
||||
},
|
||||
});
|
||||
|
||||
const data = [
|
||||
{ id: 1, text: "Carrot", type: "vegetable" },
|
||||
{ id: 2, text: "Apple", type: "fruit" },
|
||||
];
|
||||
const buffer = await fromDataToBuffer(
|
||||
data,
|
||||
undefined,
|
||||
await table.schema(),
|
||||
);
|
||||
const generatedTable = tableFromIPC(buffer);
|
||||
const vectorField = generatedTable.schema.fields.find(
|
||||
(field) => field.name === "vector",
|
||||
);
|
||||
expect(vectorField?.nullable).toBe(false);
|
||||
|
||||
await table.add(data);
|
||||
|
||||
const rows = await table.query().toArray();
|
||||
expect(rows).toHaveLength(2);
|
||||
for (const row of rows) {
|
||||
expect([...row.vector]).toEqual([1, 2, 3]);
|
||||
}
|
||||
});
|
||||
|
||||
it("should error when appending to a table with an unregistered embedding function", async () => {
|
||||
@register("mock")
|
||||
class MockEmbeddingFunction extends EmbeddingFunction<string> {
|
||||
@@ -256,6 +368,60 @@ describe("embedding functions", () => {
|
||||
expect(actual).toHaveProperty("text");
|
||||
});
|
||||
|
||||
it("should handle undefined vector field with embedding function correctly", async () => {
|
||||
@register("undefined_test")
|
||||
class MockEmbeddingFunction extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 3;
|
||||
}
|
||||
embeddingDataType(): Float {
|
||||
return new Float32();
|
||||
}
|
||||
async computeQueryEmbeddings(_data: string) {
|
||||
return [1, 2, 3];
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return Array.from({ length: data.length }).fill([
|
||||
1, 2, 3,
|
||||
]) as number[][];
|
||||
}
|
||||
}
|
||||
const func = getRegistry()
|
||||
.get<MockEmbeddingFunction>("undefined_test")!
|
||||
.create();
|
||||
const schema = new Schema([
|
||||
new Field("text", new Utf8(), true),
|
||||
new Field(
|
||||
"vector",
|
||||
new FixedSizeList(3, new Field("item", new Float32(), true)),
|
||||
true,
|
||||
),
|
||||
]);
|
||||
|
||||
const db = await connect(tmpDir.name);
|
||||
const table = await db.createEmptyTable("test_undefined", schema, {
|
||||
embeddingFunction: {
|
||||
function: func,
|
||||
sourceColumn: "text",
|
||||
vectorColumn: "vector",
|
||||
},
|
||||
});
|
||||
|
||||
// Test that undefined, null, and omitted vector fields all work
|
||||
await table.add([{ text: "test1", vector: undefined }]);
|
||||
await table.add([{ text: "test2", vector: null }]);
|
||||
await table.add([{ text: "test3" }]);
|
||||
|
||||
const rows = await table.query().toArray();
|
||||
expect(rows.length).toBe(3);
|
||||
|
||||
// All rows should have vectors computed by the embedding function
|
||||
for (const row of rows) {
|
||||
expect(row.vector).toBeDefined();
|
||||
expect(JSON.parse(JSON.stringify(row.vector))).toEqual([1, 2, 3]);
|
||||
}
|
||||
});
|
||||
|
||||
test.each([new Float16(), new Float32(), new Float64()])(
|
||||
"should be able to provide manual embeddings with multiple float datatype",
|
||||
async (floatType) => {
|
||||
@@ -373,4 +539,52 @@ describe("embedding functions", () => {
|
||||
expect(stringSchema3).toEqual(stringExpectedSchema);
|
||||
},
|
||||
);
|
||||
test("parses one function writing several vector columns", async () => {
|
||||
class MockEmbeddingFunction extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 3;
|
||||
}
|
||||
embeddingDataType(): Float {
|
||||
return new Float32();
|
||||
}
|
||||
async computeQueryEmbeddings(_data: string) {
|
||||
return [1, 2, 3];
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return Array.from({ length: data.length }).fill([
|
||||
1, 2, 3,
|
||||
]) as number[][];
|
||||
}
|
||||
}
|
||||
const registry = getRegistry();
|
||||
registry.register("multi_output_mock")(MockEmbeddingFunction);
|
||||
|
||||
// A materialized view can project one source vector column under two
|
||||
// names, so a table's configuration names the same function twice.
|
||||
const parsed = await registry.parseFunctions(
|
||||
new Map([
|
||||
[
|
||||
"embedding_functions",
|
||||
JSON.stringify([
|
||||
{
|
||||
name: "multi_output_mock",
|
||||
sourceColumn: "text",
|
||||
vectorColumn: "vector_a",
|
||||
model: {},
|
||||
},
|
||||
{
|
||||
name: "multi_output_mock",
|
||||
sourceColumn: "text",
|
||||
vectorColumn: "vector_b",
|
||||
model: {},
|
||||
},
|
||||
]),
|
||||
],
|
||||
]),
|
||||
);
|
||||
|
||||
expect(
|
||||
[...parsed.values()].map(({ vectorColumn }) => vectorColumn).sort(),
|
||||
).toEqual(["vector_a", "vector_b"]);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { execFileSync } from "node:child_process";
|
||||
import { resolve } from "node:path";
|
||||
|
||||
import type { OpenAIEmbeddingFunction } from "../lancedb/embedding/openai";
|
||||
import type { EmbeddingFunctionRegistry } from "../lancedb/embedding/registry";
|
||||
|
||||
type EmbeddingModule = typeof import("../lancedb/embedding");
|
||||
type OpenAIModule = typeof import("../lancedb/embedding/openai");
|
||||
type RegistryModule = typeof import("../lancedb/embedding/registry");
|
||||
|
||||
describe("embedding function registry", () => {
|
||||
const registries: EmbeddingFunctionRegistry[] = [];
|
||||
|
||||
afterEach(() => {
|
||||
for (const registry of registries) {
|
||||
registry.reset();
|
||||
}
|
||||
registries.length = 0;
|
||||
});
|
||||
|
||||
it("defers built-in providers until the public registry API is used", () => {
|
||||
jest.isolateModules(() => {
|
||||
const embedding = require("../lancedb/embedding") as EmbeddingModule;
|
||||
const { getRegistry: getInternalRegistry } =
|
||||
require("../lancedb/embedding/registry") as RegistryModule;
|
||||
const registry = getInternalRegistry();
|
||||
registries.push(registry);
|
||||
|
||||
expect(registry.length()).toBe(0);
|
||||
expect(embedding.getRegistry()).toBe(registry);
|
||||
expect(registry.get("openai")).toBeDefined();
|
||||
expect(registry.get("huggingface")).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
it("preserves automatic FTS search in a fresh process", () => {
|
||||
execFileSync(
|
||||
process.execPath,
|
||||
[resolve(__dirname, "fixtures", "auto_fts_search.cjs")],
|
||||
{ stdio: "pipe" },
|
||||
);
|
||||
});
|
||||
|
||||
it("shares registrations across duplicated provider module graphs", () => {
|
||||
let registeringRegistry: EmbeddingFunctionRegistry | undefined;
|
||||
let latestOpenAIConstructor: typeof OpenAIEmbeddingFunction | undefined;
|
||||
|
||||
jest.isolateModules(() => {
|
||||
require("../lancedb/embedding/openai");
|
||||
const { getRegistry } =
|
||||
require("../lancedb/embedding/registry") as RegistryModule;
|
||||
registeringRegistry = getRegistry();
|
||||
registries.push(registeringRegistry);
|
||||
expect(registeringRegistry.get("openai")).toBeDefined();
|
||||
});
|
||||
|
||||
expect(() => {
|
||||
jest.isolateModules(() => {
|
||||
const { OpenAIEmbeddingFunction } =
|
||||
require("../lancedb/embedding/openai") as OpenAIModule;
|
||||
latestOpenAIConstructor = OpenAIEmbeddingFunction;
|
||||
const { getRegistry } =
|
||||
require("../lancedb/embedding/registry") as RegistryModule;
|
||||
registries.push(getRegistry());
|
||||
});
|
||||
}).not.toThrow();
|
||||
|
||||
const previousApiKey = process.env.OPENAI_API_KEY;
|
||||
process.env.OPENAI_API_KEY = "test";
|
||||
try {
|
||||
const latestOpenAI = registeringRegistry!
|
||||
.get<OpenAIEmbeddingFunction>("openai")!
|
||||
.create();
|
||||
expect(latestOpenAI).toBeInstanceOf(latestOpenAIConstructor!);
|
||||
} finally {
|
||||
if (previousApiKey === undefined) {
|
||||
delete process.env.OPENAI_API_KEY;
|
||||
} else {
|
||||
process.env.OPENAI_API_KEY = previousApiKey;
|
||||
}
|
||||
}
|
||||
|
||||
jest.isolateModules(() => {
|
||||
const { getRegistry } =
|
||||
require("../lancedb/embedding") as EmbeddingModule;
|
||||
const publicRegistry = getRegistry();
|
||||
registries.push(publicRegistry);
|
||||
expect(publicRegistry).toBe(registeringRegistry);
|
||||
expect(publicRegistry.get("openai")).toBeDefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,33 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
const assert = require("node:assert/strict");
|
||||
const tmp = require("tmp");
|
||||
const { connect, embedding, Index } = require("../../dist");
|
||||
const { getRegistry } = require("../../dist/embedding/registry");
|
||||
|
||||
async function main() {
|
||||
assert.equal(typeof embedding.getRegistry, "function");
|
||||
assert.equal(getRegistry().length(), 0);
|
||||
assert.equal(embedding.getRegistry(), getRegistry());
|
||||
assert.equal(getRegistry().length(), 2);
|
||||
|
||||
const dir = tmp.dirSync({ unsafeCleanup: true });
|
||||
let db;
|
||||
try {
|
||||
db = await connect(dir.name);
|
||||
const table = await db.createTable("docs", [{ text: "hello world" }]);
|
||||
await table.createIndex("text", { config: Index.fts() });
|
||||
|
||||
const rows = await table.search("hello").toArray();
|
||||
assert.equal(rows[0].text, "hello world");
|
||||
} finally {
|
||||
db?.close();
|
||||
dir.removeCallback();
|
||||
}
|
||||
}
|
||||
|
||||
main().catch((error) => {
|
||||
console.error(error);
|
||||
process.exitCode = 1;
|
||||
});
|
||||
@@ -0,0 +1,3 @@
|
||||
@rem SPDX-License-Identifier: Apache-2.0
|
||||
@rem SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
@exit /b 0
|
||||
@@ -0,0 +1,181 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as tmp from "tmp";
|
||||
|
||||
import { Connection, connect } from "../lancedb";
|
||||
import {
|
||||
DEFINITION_META_KEY,
|
||||
definitionFromMetadata,
|
||||
} from "../lancedb/materialized_view";
|
||||
|
||||
describe("materialized views", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let db: Connection;
|
||||
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
db = await connect(tmpDir.name);
|
||||
await db.createTable("people", [
|
||||
{ name: "ada", age: 36 },
|
||||
{ name: "kid", age: 7 },
|
||||
{ name: "grace", age: 85 },
|
||||
]);
|
||||
});
|
||||
afterEach(() => tmpDir.removeCallback());
|
||||
|
||||
it("reads stored queries and legacy layouts", () => {
|
||||
const read = (stored: string) =>
|
||||
definitionFromMetadata(new Map([[DEFINITION_META_KEY, stored]]), "v");
|
||||
const query =
|
||||
"SELECT id, c.chunk FROM ns.docs, UNNEST(chunks) AS c WHERE id > 1";
|
||||
expect(read(`{"format":1,"query":${JSON.stringify(query)}}`).query).toBe(
|
||||
query,
|
||||
);
|
||||
|
||||
// The structured layout written before the format number reads as the
|
||||
// query it described, under either of its kind tags.
|
||||
expect(
|
||||
read(
|
||||
'{"kind":"namespaced_select","source_table":"people","source_namespace":["ns"],' +
|
||||
'"projections":[{"output":"name","expression":"`name`"},' +
|
||||
'{"output":"Shout","expression":"upper(name)"}],"filter":"age >= 18","limit":42}',
|
||||
).query,
|
||||
).toBe(
|
||||
"SELECT `name`, upper(name) AS `Shout` FROM ns.people WHERE age >= 18 LIMIT 42",
|
||||
);
|
||||
expect(read('{"kind":"select","source_table":"people"}').query).toBe(
|
||||
"SELECT * FROM people",
|
||||
);
|
||||
|
||||
// A newer writer's layout is reported, never guessed at.
|
||||
for (const newer of [
|
||||
`{"format":3,"query":${JSON.stringify(query)}}`,
|
||||
'{"kind":"select_v3","source_table":"people"}',
|
||||
]) {
|
||||
expect(() => read(newer)).toThrow(/cannot refresh/);
|
||||
}
|
||||
});
|
||||
|
||||
it("rejects a stored limit a number cannot carry", () => {
|
||||
const big = new Map([
|
||||
[
|
||||
DEFINITION_META_KEY,
|
||||
'{"kind":"select","source_table":"people","limit":9007199254740993}',
|
||||
],
|
||||
]);
|
||||
expect(() => definitionFromMetadata(big, "v")).toThrow(
|
||||
/too large to represent exactly/,
|
||||
);
|
||||
});
|
||||
|
||||
it("creates, refreshes and queries a view", async () => {
|
||||
const view = await db.createMaterializedView("adults", "people", {
|
||||
select: ["name", ["shout", "upper(name)"]],
|
||||
where: "age >= 18",
|
||||
});
|
||||
expect(view.name).toBe("adults");
|
||||
expect(await view.table().countRows()).toBe(2);
|
||||
|
||||
const rows = await view.table().query().toArray();
|
||||
expect(rows.map((r) => r.shout).sort()).toEqual(["ADA", "GRACE"]);
|
||||
});
|
||||
|
||||
it("round-trips the definition", async () => {
|
||||
await db.createMaterializedView("adults", "people", {
|
||||
where: "age >= 18",
|
||||
});
|
||||
const view = await db.openMaterializedView("adults");
|
||||
const definition = await view.definition();
|
||||
expect(definition.query).toBe(
|
||||
"SELECT name, age FROM people WHERE age >= 18",
|
||||
);
|
||||
});
|
||||
|
||||
it("fully refreshes after an append", async () => {
|
||||
const view = await db.createMaterializedView("copy", "people", {
|
||||
withNoData: true,
|
||||
});
|
||||
await view.refresh();
|
||||
|
||||
const people = await db.openTable("people");
|
||||
await people.add([{ name: "alan", age: 41 }]);
|
||||
const result = await view.refresh();
|
||||
expect(result.mode).toBe("rebuild");
|
||||
expect(Number(result.rowsWritten)).toBe(4);
|
||||
expect(await view.table().countRows()).toBe(4);
|
||||
|
||||
expect((await view.refresh()).mode).toBe("rebuild");
|
||||
});
|
||||
|
||||
it("lists views and rejects non-views", async () => {
|
||||
await db.createMaterializedView("adults", "people", {
|
||||
where: "age >= 18",
|
||||
});
|
||||
expect(await collect(db.listMaterializedViews())).toEqual(["adults"]);
|
||||
await expect(db.openMaterializedView("people")).rejects.toThrow(
|
||||
"not a materialized view",
|
||||
);
|
||||
await expect(db.dropMaterializedView("people")).rejects.toThrow(
|
||||
"not a materialized view",
|
||||
);
|
||||
|
||||
await db.dropMaterializedView("adults");
|
||||
expect(await collect(db.listMaterializedViews())).toEqual([]);
|
||||
});
|
||||
|
||||
it("returns a job when dropping a view asynchronously", async () => {
|
||||
await db.createMaterializedView("adults", "people");
|
||||
|
||||
const job = await db.dropMaterializedViewAsync("adults");
|
||||
expect(job.id).toBeNull();
|
||||
await job.wait();
|
||||
expect(await collect(db.listMaterializedViews())).toEqual([]);
|
||||
});
|
||||
|
||||
it("rejects an invalid expression at create time", async () => {
|
||||
await expect(
|
||||
db.createMaterializedView("bad", "people", {
|
||||
select: [["x", "missing + 1"]],
|
||||
}),
|
||||
).rejects.toThrow("missing");
|
||||
});
|
||||
|
||||
it("rejects invalid numeric options before creating anything", async () => {
|
||||
for (const limit of [-5, 1.5, Infinity, NaN]) {
|
||||
await expect(
|
||||
db.createMaterializedView("bad", "people", { limit }),
|
||||
).rejects.toThrow("non-negative integer");
|
||||
}
|
||||
expect(await collect(db.listMaterializedViews())).toEqual([]);
|
||||
|
||||
const view = await db.createMaterializedView("copy", "people");
|
||||
for (const sourceVersion of [-1, 1.5, Infinity, NaN]) {
|
||||
await expect(view.refresh({ sourceVersion })).rejects.toThrow(
|
||||
"non-negative integer",
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
it("quotes bare select names", async () => {
|
||||
await db.createTable("odd_names", [{ "order item": "widget" }]);
|
||||
const view = await db.createMaterializedView("quoted", "odd_names", {
|
||||
select: ["order item"],
|
||||
withNoData: true,
|
||||
});
|
||||
const result = await view.refresh();
|
||||
expect(Number(result.rowsWritten)).toBe(1);
|
||||
});
|
||||
|
||||
it("does not require stable row ids on the source", async () => {
|
||||
await db.createTable("plain", [{ x: 1 }]);
|
||||
const view = await db.createMaterializedView("v", "plain");
|
||||
expect(await view.table().countRows()).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
async function collect<T>(items: AsyncIterable<T>): Promise<T[]> {
|
||||
const result: T[] = [];
|
||||
for await (const item of items) result.push(item);
|
||||
return result;
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as fs from "fs";
|
||||
import * as http from "http";
|
||||
import { execFileSync } from "node:child_process";
|
||||
import * as os from "os";
|
||||
import * as path from "path";
|
||||
import { OAuthConfig, OAuthFlowType, OAuthSession } from "../lancedb/oauth";
|
||||
|
||||
function tempCacheDir(): string {
|
||||
return fs.mkdtempSync(path.join(os.tmpdir(), "lancedb-oauth-cache-"));
|
||||
}
|
||||
|
||||
function deviceConfig(issuerUrl: string, cacheDir: string): OAuthConfig {
|
||||
return {
|
||||
issuerUrl,
|
||||
clientId: "client-id",
|
||||
scopes: ["openid"],
|
||||
flow: OAuthFlowType.DeviceCode,
|
||||
tokenCache: { cacheDir },
|
||||
};
|
||||
}
|
||||
|
||||
describe("OAuthSession", () => {
|
||||
beforeAll(() => {
|
||||
// Child processes inherit the real environment, just as Rust reads it.
|
||||
// Fail before login if the browser override only exists in Jest's sandbox.
|
||||
const browser = execFileSync(
|
||||
process.execPath,
|
||||
["-p", "process.env.LANCEDB_OAUTH_BROWSER ?? ''"],
|
||||
{ encoding: "utf8" },
|
||||
).trim();
|
||||
expect(browser).not.toBe("");
|
||||
expect(browser).toBe(process.env.LANCEDB_OAUTH_BROWSER);
|
||||
});
|
||||
|
||||
it("reports an absent session and logout is idempotent", async () => {
|
||||
const cacheDir = tempCacheDir();
|
||||
const session = new OAuthSession(
|
||||
deviceConfig("https://issuer.example.com", cacheDir),
|
||||
);
|
||||
|
||||
const status = await session.status();
|
||||
expect(status.refreshable).toBe(false);
|
||||
expect(status.issuerUrl).toBe("https://issuer.example.com");
|
||||
expect(status.clientId).toBe("client-id");
|
||||
expect(status.scopes).toEqual(["openid"]);
|
||||
expect(status.flow).toBe("device_code");
|
||||
expect(status.obtainedAt).toBeUndefined();
|
||||
|
||||
const logout = await session.logout();
|
||||
expect(logout.removed).toBe(false);
|
||||
});
|
||||
|
||||
it("requires token cache options", () => {
|
||||
const config: OAuthConfig = {
|
||||
issuerUrl: "https://issuer.example.com",
|
||||
clientId: "client-id",
|
||||
scopes: ["openid"],
|
||||
flow: OAuthFlowType.DeviceCode,
|
||||
};
|
||||
expect(() => new OAuthSession(config)).toThrow(/token/);
|
||||
});
|
||||
|
||||
it("rejects azure managed identity persistence", () => {
|
||||
const config: OAuthConfig = {
|
||||
issuerUrl: "https://login.microsoftonline.com/tenant/v2.0",
|
||||
clientId: "app-id",
|
||||
scopes: ["api://app/.default"],
|
||||
flow: OAuthFlowType.AzureManagedIdentity,
|
||||
tokenCache: { cacheDir: tempCacheDir() },
|
||||
};
|
||||
expect(() => new OAuthSession(config)).toThrow(/AzureManagedIdentity/);
|
||||
});
|
||||
|
||||
it.each([
|
||||
{},
|
||||
{
|
||||
resource: "https://api.example.com/a?x=1&y=two",
|
||||
audience: "audience + & / ü",
|
||||
},
|
||||
])(
|
||||
"logs in via device flow with target %j, caches, and logs out",
|
||||
async (target) => {
|
||||
const server = new MockIdp();
|
||||
await server.start();
|
||||
try {
|
||||
const cacheDir = tempCacheDir();
|
||||
const issuerUrl = server.issuerUrl();
|
||||
|
||||
const config = { ...deviceConfig(issuerUrl, cacheDir), ...target };
|
||||
const session = new OAuthSession(config);
|
||||
const status = await session.login();
|
||||
expect(status.refreshable).toBe(true);
|
||||
expect(status.resource).toBe(config.resource);
|
||||
expect(status.audience).toBe(config.audience);
|
||||
expect(status.obtainedAt).toBeGreaterThan(0);
|
||||
expect(server.state.deviceAuthorizations).toBe(1);
|
||||
|
||||
// An independent session (a fresh "process") sees the cached login.
|
||||
const other = new OAuthSession(config);
|
||||
const cached = await other.status();
|
||||
expect(cached.refreshable).toBe(true);
|
||||
|
||||
const logout = await other.logout();
|
||||
expect(logout.removed).toBe(true);
|
||||
const again = await session.logout();
|
||||
expect(again.removed).toBe(false);
|
||||
expect((await session.status()).refreshable).toBe(false);
|
||||
|
||||
// Only the initial login used the interactive device flow.
|
||||
expect(server.state.deviceAuthorizations).toBe(1);
|
||||
expect(server.state.refreshGrants).toBe(0);
|
||||
expect(server.requests).toHaveLength(2);
|
||||
for (const params of server.requests) {
|
||||
expect(params.getAll("resource")).toEqual(
|
||||
config.resource === undefined ? [] : [config.resource],
|
||||
);
|
||||
expect(params.getAll("audience")).toEqual(
|
||||
config.audience === undefined ? [] : [config.audience],
|
||||
);
|
||||
}
|
||||
} finally {
|
||||
server.close();
|
||||
}
|
||||
},
|
||||
15000,
|
||||
);
|
||||
});
|
||||
|
||||
/** Mock IdP with discovery, device authorization, and rotating refresh. */
|
||||
class MockIdp {
|
||||
readonly requests: URLSearchParams[] = [];
|
||||
readonly state = {
|
||||
deviceAuthorizations: 0,
|
||||
refreshGrants: 0,
|
||||
accessTokensIssued: 0,
|
||||
currentRefresh: null as string | null,
|
||||
};
|
||||
private server?: http.Server;
|
||||
private port = 0;
|
||||
|
||||
issuerUrl(): string {
|
||||
return `http://127.0.0.1:${this.port}`;
|
||||
}
|
||||
|
||||
async start(): Promise<void> {
|
||||
const server = http.createServer((req, res) => {
|
||||
const chunks: Buffer[] = [];
|
||||
req.on("data", (chunk) => chunks.push(chunk));
|
||||
req.on("end", () => {
|
||||
const body = Buffer.concat(chunks).toString();
|
||||
const params = new URLSearchParams(body);
|
||||
this.handle(req.url ?? "", params, res);
|
||||
});
|
||||
});
|
||||
await new Promise<void>((resolve) => {
|
||||
server.listen(0, "127.0.0.1", () => resolve());
|
||||
});
|
||||
const address = server.address();
|
||||
if (address && typeof address === "object") {
|
||||
this.port = address.port;
|
||||
}
|
||||
this.server = server;
|
||||
}
|
||||
|
||||
private handle(
|
||||
url: string,
|
||||
params: URLSearchParams,
|
||||
res: http.ServerResponse,
|
||||
): void {
|
||||
const respond = (status: number, payload: unknown): void => {
|
||||
const body = JSON.stringify(payload);
|
||||
res.writeHead(status, {
|
||||
"Content-Type": "application/json",
|
||||
"Content-Length": Buffer.byteLength(body),
|
||||
});
|
||||
res.end(body);
|
||||
};
|
||||
|
||||
if (url === "/.well-known/openid-configuration") {
|
||||
respond(200, {
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
token_endpoint: `${this.issuerUrl()}/token`,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
device_authorization_endpoint: `${this.issuerUrl()}/device`,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (url === "/device" || url === "/token") {
|
||||
this.requests.push(params);
|
||||
}
|
||||
|
||||
if (url === "/device") {
|
||||
this.state.deviceAuthorizations += 1;
|
||||
respond(200, {
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
device_code: "device-code",
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
user_code: "ABCD-EFGH",
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
verification_uri: `${this.issuerUrl()}/verify`,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
expires_in: 60,
|
||||
interval: 1,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (url === "/token") {
|
||||
const grantType = params.get("grant_type") ?? "";
|
||||
if (grantType === "refresh_token") {
|
||||
this.state.refreshGrants += 1;
|
||||
if (params.get("refresh_token") !== this.state.currentRefresh) {
|
||||
respond(400, { error: "invalid_grant" });
|
||||
return;
|
||||
}
|
||||
} else if (!grantType.includes("device_code")) {
|
||||
respond(400, { error: "unsupported_grant_type" });
|
||||
return;
|
||||
}
|
||||
this.state.accessTokensIssued += 1;
|
||||
const number = this.state.accessTokensIssued;
|
||||
const refresh = `refresh-${number}`;
|
||||
this.state.currentRefresh = refresh;
|
||||
respond(200, {
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
access_token: `access-${number}`,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
refresh_token: refresh,
|
||||
// biome-ignore lint/style/useNamingConvention: OAuth wire format
|
||||
expires_in: 3600,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
respond(404, {});
|
||||
}
|
||||
|
||||
close(): void {
|
||||
this.server?.close();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import {
|
||||
MeterProvider,
|
||||
type MetricData,
|
||||
MetricReader,
|
||||
} from "@opentelemetry/sdk-metrics";
|
||||
import * as tmp from "tmp";
|
||||
import { connect, instrumentLanceDbMetrics } from "../lancedb";
|
||||
// snapshotLancedbMetrics is internal plumbing (not part of the public API), so
|
||||
// it is imported from the native module rather than the package entry point.
|
||||
import { snapshotLancedbMetrics } from "../lancedb/native";
|
||||
|
||||
// The metrics recorder is process-global and installed once, so the whole
|
||||
// bridge is exercised in a single test to avoid cross-test global-state coupling.
|
||||
|
||||
// A minimal pull-based reader whose `collect()` we drive directly, invoking the
|
||||
// observable-instrument callbacks. `@opentelemetry/sdk-metrics` ships no
|
||||
// in-memory reader, so we subclass the abstract base.
|
||||
class TestMetricReader extends MetricReader {
|
||||
protected async onForceFlush(): Promise<void> {
|
||||
// no-op: collection is driven directly via collect()
|
||||
}
|
||||
protected async onShutdown(): Promise<void> {
|
||||
// no-op: nothing to release
|
||||
}
|
||||
}
|
||||
|
||||
async function metricsByName(
|
||||
reader: TestMetricReader,
|
||||
): Promise<Map<string, MetricData>> {
|
||||
const collected = await reader.collect();
|
||||
const result = new Map<string, MetricData>();
|
||||
for (const scope of collected.resourceMetrics.scopeMetrics) {
|
||||
for (const metric of scope.metrics) {
|
||||
result.set(metric.descriptor.name, metric);
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
describe("OpenTelemetry metrics bridge", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
beforeEach(() => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
});
|
||||
afterEach(() => tmpDir.removeCallback());
|
||||
|
||||
it("snapshot is safe to call regardless of install state", () => {
|
||||
expect(Array.isArray(snapshotLancedbMetrics())).toBe(true);
|
||||
});
|
||||
|
||||
it("exports object store metrics via observable instruments", async () => {
|
||||
const reader = new TestMetricReader();
|
||||
const provider = new MeterProvider({ readers: [reader] });
|
||||
expect(instrumentLanceDbMetrics(provider)).toBe(true);
|
||||
|
||||
// Generate object store activity on the local filesystem (scheme "file").
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = Array.from({ length: 256 }, (_, i) => ({ id: i }));
|
||||
const table = await db.createTable("t", data);
|
||||
expect(await table.countRows()).toBe(256);
|
||||
|
||||
const metrics = await metricsByName(reader);
|
||||
|
||||
const requests = metrics.get("lance_object_store_requests_total");
|
||||
expect(requests).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const requestPoints = (requests!.dataPoints as any[]) ?? [];
|
||||
expect(requestPoints.length).toBeGreaterThan(0);
|
||||
for (const p of requestPoints) {
|
||||
// Labelled by `operation` and `base` (the store scheme by default).
|
||||
expect(p.attributes).toHaveProperty("base");
|
||||
expect(p.attributes).toHaveProperty("operation");
|
||||
}
|
||||
const totalRequests = requestPoints.reduce((acc, p) => acc + p.value, 0);
|
||||
expect(totalRequests).toBeGreaterThan(0);
|
||||
|
||||
// Histograms are decomposed into bucket / count / sum observable counters.
|
||||
const bucket = metrics.get(
|
||||
"lance_object_store_request_duration_seconds_bucket",
|
||||
);
|
||||
expect(bucket).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const bucketPoints = (bucket!.dataPoints as any[]) ?? [];
|
||||
expect(bucketPoints.length).toBeGreaterThan(0);
|
||||
expect(bucketPoints.every((p) => "le" in p.attributes)).toBe(true);
|
||||
// The implicit +Inf bucket must be present.
|
||||
expect(bucketPoints.some((p) => p.attributes.le === "+Inf")).toBe(true);
|
||||
|
||||
const count = metrics.get(
|
||||
"lance_object_store_request_duration_seconds_count",
|
||||
);
|
||||
expect(count).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const countPoints = (count!.dataPoints as any[]) ?? [];
|
||||
expect(countPoints.reduce((acc, p) => acc + p.value, 0)).toBeGreaterThan(0);
|
||||
|
||||
const sum = metrics.get("lance_object_store_request_duration_seconds_sum");
|
||||
expect(sum).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const sumPoints = (sum!.dataPoints as any[]) ?? [];
|
||||
expect(sumPoints.reduce((acc, p) => acc + p.value, 0)).toBeGreaterThan(0);
|
||||
|
||||
// Unit handling: only `_sum` keeps the histogram's unit (seconds); `_bucket`
|
||||
// and `_count` observe cumulative counts and are unitless.
|
||||
expect(sum!.descriptor.unit).toBe("s");
|
||||
expect(bucket!.descriptor.unit).toBe("");
|
||||
expect(count!.descriptor.unit).toBe("");
|
||||
|
||||
await provider.shutdown();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,14 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import packageJson = require("../package.json");
|
||||
|
||||
describe("package metadata", () => {
|
||||
it("requires Node.js type declarations compatible with the runtime", () => {
|
||||
expect(packageJson.engines.node).toBe(">= 22");
|
||||
expect(packageJson.peerDependencies["@types/node"]).toBe(">=22");
|
||||
expect(packageJson.peerDependenciesMeta["@types/node"]).toEqual({
|
||||
optional: true,
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,371 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as tmp from "tmp";
|
||||
import { Table, connect, permutationBuilder } from "../lancedb";
|
||||
import { makeArrowTable } from "../lancedb/arrow";
|
||||
|
||||
describe("PermutationBuilder", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let table: Table;
|
||||
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
const db = await connect(tmpDir.name);
|
||||
|
||||
// Create test data
|
||||
const data = makeArrowTable(
|
||||
[
|
||||
{ id: 1, value: 10 },
|
||||
{ id: 2, value: 20 },
|
||||
{ id: 3, value: 30 },
|
||||
{ id: 4, value: 40 },
|
||||
{ id: 5, value: 50 },
|
||||
{ id: 6, value: 60 },
|
||||
{ id: 7, value: 70 },
|
||||
{ id: 8, value: 80 },
|
||||
{ id: 9, value: 90 },
|
||||
{ id: 10, value: 100 },
|
||||
],
|
||||
{ vectorColumns: {} },
|
||||
);
|
||||
|
||||
table = await db.createTable("test_table", data);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
tmpDir.removeCallback();
|
||||
});
|
||||
|
||||
test("should create permutation builder", () => {
|
||||
const builder = permutationBuilder(table);
|
||||
expect(builder).toBeDefined();
|
||||
});
|
||||
|
||||
test("should execute basic permutation", async () => {
|
||||
const builder = permutationBuilder(table);
|
||||
const permutationTable = await builder.execute();
|
||||
|
||||
expect(permutationTable).toBeDefined();
|
||||
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
});
|
||||
|
||||
test("should create permutation with random splits", async () => {
|
||||
const builder = permutationBuilder(table).splitRandom({
|
||||
ratios: [1.0],
|
||||
seed: 42,
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
});
|
||||
|
||||
test("should create permutation with percentage splits", async () => {
|
||||
const builder = permutationBuilder(table).splitRandom({
|
||||
ratios: [0.3, 0.7],
|
||||
seed: 42,
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
|
||||
// Check split distribution
|
||||
const split0Count = await permutationTable.countRows("split_id = 0");
|
||||
const split1Count = await permutationTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBeGreaterThan(0);
|
||||
expect(split1Count).toBeGreaterThan(0);
|
||||
expect(split0Count + split1Count).toBe(10);
|
||||
});
|
||||
|
||||
test("should create permutation with count splits", async () => {
|
||||
const builder = permutationBuilder(table).splitRandom({
|
||||
counts: [3, 7],
|
||||
seed: 42,
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
|
||||
// Check split distribution
|
||||
const split0Count = await permutationTable.countRows("split_id = 0");
|
||||
const split1Count = await permutationTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBe(3);
|
||||
expect(split1Count).toBe(7);
|
||||
});
|
||||
|
||||
test("should create permutation with hash splits", async () => {
|
||||
const builder = permutationBuilder(table).splitHash({
|
||||
columns: ["id"],
|
||||
splitWeights: [50, 50],
|
||||
discardWeight: 0,
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
|
||||
// Check that splits exist
|
||||
const split0Count = await permutationTable.countRows("split_id = 0");
|
||||
const split1Count = await permutationTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBeGreaterThan(0);
|
||||
expect(split1Count).toBeGreaterThan(0);
|
||||
expect(split0Count + split1Count).toBe(10);
|
||||
});
|
||||
|
||||
test("should create permutation with sequential splits", async () => {
|
||||
const builder = permutationBuilder(table).splitSequential({
|
||||
ratios: [0.5, 0.5],
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
|
||||
// Check split distribution - sequential should give exactly 5 and 5
|
||||
const split0Count = await permutationTable.countRows("split_id = 0");
|
||||
const split1Count = await permutationTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBe(5);
|
||||
expect(split1Count).toBe(5);
|
||||
});
|
||||
|
||||
test("should create permutation with calculated splits", async () => {
|
||||
const builder = permutationBuilder(table).splitCalculated({
|
||||
calculation: "id % 2",
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
|
||||
// Check split distribution
|
||||
const split0Count = await permutationTable.countRows("split_id = 0");
|
||||
const split1Count = await permutationTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBeGreaterThan(0);
|
||||
expect(split1Count).toBeGreaterThan(0);
|
||||
expect(split0Count + split1Count).toBe(10);
|
||||
});
|
||||
|
||||
test("should create permutation with shuffle", async () => {
|
||||
const builder = permutationBuilder(table).shuffle({
|
||||
seed: 42,
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
});
|
||||
|
||||
test("should create permutation with shuffle and clump size", async () => {
|
||||
const builder = permutationBuilder(table).shuffle({
|
||||
seed: 42,
|
||||
clumpSize: 2,
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
});
|
||||
|
||||
test("should create permutation with filter", async () => {
|
||||
const builder = permutationBuilder(table).filter("value > 50");
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(5); // Values 60, 70, 80, 90, 100
|
||||
});
|
||||
|
||||
test("should chain multiple operations", async () => {
|
||||
const builder = permutationBuilder(table)
|
||||
.filter("value <= 80")
|
||||
.splitRandom({ ratios: [0.5, 0.5], seed: 42 })
|
||||
.shuffle({ seed: 123 });
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(8); // Values 10, 20, 30, 40, 50, 60, 70, 80
|
||||
|
||||
// Check split distribution
|
||||
const split0Count = await permutationTable.countRows("split_id = 0");
|
||||
const split1Count = await permutationTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBeGreaterThan(0);
|
||||
expect(split1Count).toBeGreaterThan(0);
|
||||
expect(split0Count + split1Count).toBe(8);
|
||||
});
|
||||
|
||||
test("should throw error for invalid split arguments", () => {
|
||||
const builder = permutationBuilder(table);
|
||||
|
||||
// Test no arguments provided
|
||||
expect(() => builder.splitRandom({})).toThrow(
|
||||
"Exactly one of 'ratios', 'counts', or 'fixed' must be provided",
|
||||
);
|
||||
|
||||
// Test multiple arguments provided
|
||||
expect(() =>
|
||||
builder.splitRandom({ ratios: [0.5, 0.5], counts: [3, 7], seed: 42 }),
|
||||
).toThrow("Exactly one of 'ratios', 'counts', or 'fixed' must be provided");
|
||||
});
|
||||
|
||||
test("should throw error when builder is consumed", async () => {
|
||||
const builder = permutationBuilder(table);
|
||||
|
||||
// Execute once
|
||||
await builder.execute();
|
||||
|
||||
// Should throw error on second execution
|
||||
await expect(builder.execute()).rejects.toThrow("Builder already consumed");
|
||||
});
|
||||
|
||||
test("should accept custom split names with random splits", async () => {
|
||||
const builder = permutationBuilder(table).splitRandom({
|
||||
ratios: [0.3, 0.7],
|
||||
seed: 42,
|
||||
splitNames: ["train", "test"],
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
|
||||
// Split names are provided but split_id is still numeric (0, 1, etc.)
|
||||
// The names are metadata that can be used by higher-level APIs
|
||||
const split0Count = await permutationTable.countRows("split_id = 0");
|
||||
const split1Count = await permutationTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBeGreaterThan(0);
|
||||
expect(split1Count).toBeGreaterThan(0);
|
||||
expect(split0Count + split1Count).toBe(10);
|
||||
});
|
||||
|
||||
test("should accept custom split names with hash splits", async () => {
|
||||
const builder = permutationBuilder(table).splitHash({
|
||||
columns: ["id"],
|
||||
splitWeights: [50, 50],
|
||||
discardWeight: 0,
|
||||
splitNames: ["set_a", "set_b"],
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
|
||||
// Split names are provided but split_id is still numeric
|
||||
const split0Count = await permutationTable.countRows("split_id = 0");
|
||||
const split1Count = await permutationTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBeGreaterThan(0);
|
||||
expect(split1Count).toBeGreaterThan(0);
|
||||
expect(split0Count + split1Count).toBe(10);
|
||||
});
|
||||
|
||||
test("should accept custom split names with sequential splits", async () => {
|
||||
const builder = permutationBuilder(table).splitSequential({
|
||||
ratios: [0.5, 0.5],
|
||||
splitNames: ["first", "second"],
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
|
||||
// Split names are provided but split_id is still numeric
|
||||
const split0Count = await permutationTable.countRows("split_id = 0");
|
||||
const split1Count = await permutationTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBe(5);
|
||||
expect(split1Count).toBe(5);
|
||||
});
|
||||
|
||||
test("should accept custom split names with calculated splits", async () => {
|
||||
const builder = permutationBuilder(table).splitCalculated({
|
||||
calculation: "id % 2",
|
||||
splitNames: ["even", "odd"],
|
||||
});
|
||||
|
||||
const permutationTable = await builder.execute();
|
||||
const rowCount = await permutationTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
|
||||
// Split names are provided but split_id is still numeric
|
||||
const split0Count = await permutationTable.countRows("split_id = 0");
|
||||
const split1Count = await permutationTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBeGreaterThan(0);
|
||||
expect(split1Count).toBeGreaterThan(0);
|
||||
expect(split0Count + split1Count).toBe(10);
|
||||
});
|
||||
|
||||
test("should persist permutation to a new table", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const builder = permutationBuilder(table)
|
||||
.splitRandom({
|
||||
ratios: [0.7, 0.3],
|
||||
seed: 42,
|
||||
splitNames: ["train", "validation"],
|
||||
})
|
||||
.persist(db, "my_permutation");
|
||||
|
||||
// Execute the builder which will persist the table
|
||||
const permutationTable = await builder.execute();
|
||||
|
||||
// Verify the persisted table exists and can be opened
|
||||
const persistedTable = await db.openTable("my_permutation");
|
||||
expect(persistedTable).toBeDefined();
|
||||
|
||||
// Verify the persisted table has the correct number of rows
|
||||
const rowCount = await persistedTable.countRows();
|
||||
expect(rowCount).toBe(10);
|
||||
|
||||
// Verify splits exist (numeric split_id values)
|
||||
const split0Count = await persistedTable.countRows("split_id = 0");
|
||||
const split1Count = await persistedTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBeGreaterThan(0);
|
||||
expect(split1Count).toBeGreaterThan(0);
|
||||
expect(split0Count + split1Count).toBe(10);
|
||||
|
||||
// Verify the table returned by execute is the same as the persisted one
|
||||
const executedRowCount = await permutationTable.countRows();
|
||||
expect(executedRowCount).toBe(10);
|
||||
});
|
||||
|
||||
test("should persist permutation with multiple operations", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const builder = permutationBuilder(table)
|
||||
.filter("value > 30")
|
||||
.splitRandom({ ratios: [0.5, 0.5], seed: 123, splitNames: ["a", "b"] })
|
||||
.shuffle({ seed: 456 })
|
||||
.persist(db, "filtered_permutation");
|
||||
|
||||
// Execute the builder
|
||||
const permutationTable = await builder.execute();
|
||||
|
||||
// Verify the persisted table
|
||||
const persistedTable = await db.openTable("filtered_permutation");
|
||||
const rowCount = await persistedTable.countRows();
|
||||
expect(rowCount).toBe(7); // Values 40, 50, 60, 70, 80, 90, 100
|
||||
|
||||
// Verify splits exist (numeric split_id values)
|
||||
const split0Count = await persistedTable.countRows("split_id = 0");
|
||||
const split1Count = await persistedTable.countRows("split_id = 1");
|
||||
|
||||
expect(split0Count).toBeGreaterThan(0);
|
||||
expect(split1Count).toBeGreaterThan(0);
|
||||
expect(split0Count + split1Count).toBe(7);
|
||||
|
||||
// Verify the executed table matches
|
||||
const executedRowCount = await permutationTable.countRows();
|
||||
expect(executedRowCount).toBe(7);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,696 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as tmp from "tmp";
|
||||
|
||||
import { type Table, connect } from "../lancedb";
|
||||
import {
|
||||
Field,
|
||||
FixedSizeList,
|
||||
Float32,
|
||||
Int64,
|
||||
Schema,
|
||||
Utf8,
|
||||
makeArrowTable,
|
||||
} from "../lancedb/arrow";
|
||||
import { Index } from "../lancedb/indices";
|
||||
|
||||
describe("Query outputSchema", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let table: Table;
|
||||
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
const db = await connect(tmpDir.name);
|
||||
|
||||
// Create table with explicit schema to ensure proper types
|
||||
const schema = new Schema([
|
||||
new Field("a", new Int64(), true),
|
||||
new Field("text", new Utf8(), true),
|
||||
new Field(
|
||||
"vec",
|
||||
new FixedSizeList(2, new Field("item", new Float32())),
|
||||
true,
|
||||
),
|
||||
]);
|
||||
|
||||
const data = makeArrowTable(
|
||||
[
|
||||
{ a: 1n, text: "foo", vec: [1, 2] },
|
||||
{ a: 2n, text: "bar", vec: [3, 4] },
|
||||
{ a: 3n, text: "baz", vec: [5, 6] },
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
table = await db.createTable("test", data);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
tmpDir.removeCallback();
|
||||
});
|
||||
|
||||
it("should return schema for plain query", async () => {
|
||||
const schema = await table.query().outputSchema();
|
||||
|
||||
expect(schema.fields.length).toBe(3);
|
||||
expect(schema.fields.map((f) => f.name)).toEqual(["a", "text", "vec"]);
|
||||
expect(schema.fields[0].type.toString()).toBe("Int64");
|
||||
expect(schema.fields[1].type.toString()).toBe("Utf8");
|
||||
});
|
||||
|
||||
it("should return schema with dynamic projection", async () => {
|
||||
const schema = await table.query().select({ bl: "a * 2" }).outputSchema();
|
||||
|
||||
expect(schema.fields.length).toBe(1);
|
||||
expect(schema.fields[0].name).toBe("bl");
|
||||
expect(schema.fields[0].type.toString()).toBe("Int64");
|
||||
});
|
||||
|
||||
it("should return schema for vector search with _distance column", async () => {
|
||||
const schema = await table
|
||||
.vectorSearch([1, 2])
|
||||
.select(["a"])
|
||||
.outputSchema();
|
||||
|
||||
expect(schema.fields.length).toBe(2);
|
||||
expect(schema.fields.map((f) => f.name)).toEqual(["a", "_distance"]);
|
||||
expect(schema.fields[0].type.toString()).toBe("Int64");
|
||||
expect(schema.fields[1].type.toString()).toBe("Float32");
|
||||
});
|
||||
|
||||
it("should return schema for FTS search", async () => {
|
||||
await table.createIndex("text", { config: Index.fts() });
|
||||
|
||||
const schema = await table
|
||||
.search("foo", "fts")
|
||||
.select(["a"])
|
||||
.outputSchema();
|
||||
|
||||
// FTS search includes _score column in addition to selected columns
|
||||
expect(schema.fields.length).toBe(2);
|
||||
expect(schema.fields.map((f) => f.name)).toContain("a");
|
||||
expect(schema.fields.map((f) => f.name)).toContain("_score");
|
||||
const aField = schema.fields.find((f) => f.name === "a");
|
||||
expect(aField?.type.toString()).toBe("Int64");
|
||||
});
|
||||
|
||||
it.each([false, true])(
|
||||
"should honor vector scoring projection (explicit score: %s)",
|
||||
async (includeScore) => {
|
||||
const columns = includeScore ? ["a", "_distance"] : ["a"];
|
||||
const query = table
|
||||
.vectorSearch([1, 2])
|
||||
.select(columns)
|
||||
.disableScoringAutoprojection();
|
||||
|
||||
const results = await query.toArrow();
|
||||
expect(results.numRows).toBe(3);
|
||||
expect(results.schema.fields.map((field) => field.name)).toEqual(columns);
|
||||
expect(
|
||||
(await query.outputSchema()).fields.map((field) => field.name),
|
||||
).toEqual(columns);
|
||||
},
|
||||
);
|
||||
|
||||
it.each([false, true])(
|
||||
"should honor FTS scoring projection (explicit score: %s)",
|
||||
async (includeScore) => {
|
||||
await table.createIndex("text", { config: Index.fts() });
|
||||
const columns = includeScore ? ["a", "_score"] : ["a"];
|
||||
const query = table
|
||||
.query()
|
||||
.fullTextSearch("foo")
|
||||
.select(columns)
|
||||
.disableScoringAutoprojection();
|
||||
|
||||
const results = await query.toArrow();
|
||||
expect(results.numRows).toBe(1);
|
||||
expect(results.schema.fields.map((field) => field.name)).toEqual(columns);
|
||||
expect(
|
||||
(await query.outputSchema()).fields.map((field) => field.name),
|
||||
).toEqual(columns);
|
||||
},
|
||||
);
|
||||
|
||||
it("should preserve the scoring option when converting to a vector query", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.disableScoringAutoprojection()
|
||||
.nearestTo([1, 2])
|
||||
.select({ doubled: "a * 2" })
|
||||
.toArrow();
|
||||
|
||||
expect(results.numRows).toBe(3);
|
||||
expect(results.schema.fields.map((field) => field.name)).toEqual([
|
||||
"doubled",
|
||||
]);
|
||||
});
|
||||
|
||||
it("should apply the scoring option to deferred auto queries", async () => {
|
||||
await table.createIndex("text", { config: Index.fts() });
|
||||
const results = await table
|
||||
.search("foo")
|
||||
.select(["a"])
|
||||
.disableScoringAutoprojection()
|
||||
.toArrow();
|
||||
|
||||
expect(results.numRows).toBe(1);
|
||||
expect(results.schema.fields.map((field) => field.name)).toEqual(["a"]);
|
||||
});
|
||||
|
||||
it("should return schema for take query", async () => {
|
||||
const schema = await table.takeOffsets([0]).select(["text"]).outputSchema();
|
||||
|
||||
expect(schema.fields.length).toBe(1);
|
||||
expect(schema.fields[0].name).toBe("text");
|
||||
expect(schema.fields[0].type.toString()).toBe("Utf8");
|
||||
});
|
||||
|
||||
it("should return full schema when no select is specified", async () => {
|
||||
const schema = await table.query().outputSchema();
|
||||
|
||||
// Should return all columns
|
||||
expect(schema.fields.length).toBe(3);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Query empty results", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let table: Table;
|
||||
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
const db = await connect(tmpDir.name);
|
||||
const schema = new Schema(
|
||||
[
|
||||
new Field("id", new Int64(), false),
|
||||
new Field(
|
||||
"text",
|
||||
new Utf8(),
|
||||
true,
|
||||
new Map([["description", "searchable text"]]),
|
||||
),
|
||||
new Field(
|
||||
"vector",
|
||||
new FixedSizeList(2, new Field("item", new Float32())),
|
||||
true,
|
||||
),
|
||||
],
|
||||
new Map([["source", "empty-query-regression"]]),
|
||||
);
|
||||
table = await db.createTable(
|
||||
"test",
|
||||
makeArrowTable([{ id: 1n, text: "hello", vector: [1, 2] }], { schema }),
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
tmpDir.removeCallback();
|
||||
});
|
||||
|
||||
it.each([
|
||||
{
|
||||
name: "filtered scan",
|
||||
query: (table: Table) => table.query().where("id < 0"),
|
||||
fields: ["id", "text", "vector"],
|
||||
},
|
||||
{
|
||||
name: "zero limit",
|
||||
query: (table: Table) => table.query().limit(0),
|
||||
fields: ["id", "text", "vector"],
|
||||
},
|
||||
{
|
||||
name: "column projection",
|
||||
query: (table: Table) =>
|
||||
table.query().where("id < 0").select(["text", "id"]),
|
||||
fields: ["text", "id"],
|
||||
},
|
||||
{
|
||||
name: "dynamic projection",
|
||||
query: (table: Table) =>
|
||||
table.query().where("id < 0").select({ doubled: "id * 2" }),
|
||||
fields: ["doubled"],
|
||||
},
|
||||
{
|
||||
name: "filtered vector search",
|
||||
query: (table: Table) =>
|
||||
table.search([0.5, 0.5]).where("id < 0").limit(3),
|
||||
fields: ["id", "text", "vector", "_distance"],
|
||||
},
|
||||
{
|
||||
name: "vector projection with row id",
|
||||
query: (table: Table) =>
|
||||
table
|
||||
.vectorSearch([0.5, 0.5])
|
||||
.where("id < 0")
|
||||
.select(["id"])
|
||||
.withRowId(),
|
||||
fields: ["id", "_distance", "_rowid"],
|
||||
},
|
||||
{
|
||||
name: "fast search without an index",
|
||||
query: (table: Table) => table.search([0.5, 0.5]).fastSearch(),
|
||||
fields: ["id", "text", "vector", "_distance"],
|
||||
},
|
||||
{
|
||||
name: "take with no offsets",
|
||||
query: (table: Table) => table.takeOffsets([]).select(["text"]),
|
||||
fields: ["text"],
|
||||
},
|
||||
])("preserves the output schema for $name", async ({ query, fields }) => {
|
||||
const result = await query(table).toArrow({ maxBatchLength: 1 });
|
||||
|
||||
expect(result.numRows).toBe(0);
|
||||
expect(result.schema.fields.map((field) => field.name)).toEqual(fields);
|
||||
for (const name of fields) {
|
||||
expect(result.getChild(name)).not.toBeNull();
|
||||
expect(result.getChild(name)?.length).toBe(0);
|
||||
}
|
||||
});
|
||||
|
||||
it("preserves types, nullability, and metadata", async () => {
|
||||
const result = await table.query().where("id < 0").toArrow();
|
||||
const schema = await table.schema();
|
||||
|
||||
expect(result.schema).toEqual(schema);
|
||||
expect(result.getChild("id")?.type.toString()).toBe("Int64");
|
||||
expect(result.schema.fields[0].nullable).toBe(false);
|
||||
expect(result.schema.fields[1].metadata.get("description")).toBe(
|
||||
"searchable text",
|
||||
);
|
||||
expect(result.schema.metadata.get("source")).toBe("empty-query-regression");
|
||||
});
|
||||
|
||||
it("preserves the score column for full-text search with no matches", async () => {
|
||||
await table.createIndex("text", { config: Index.fts() });
|
||||
const result = await table
|
||||
.search("missing", "fts")
|
||||
.select(["id"])
|
||||
.toArrow();
|
||||
|
||||
expect(result.numRows).toBe(0);
|
||||
expect(result.schema.fields.map((field) => field.name)).toEqual([
|
||||
"id",
|
||||
"_score",
|
||||
]);
|
||||
expect(result.getChild("_score")?.type.toString()).toBe("Float32");
|
||||
});
|
||||
});
|
||||
|
||||
describe("Search pagination", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let table: Table;
|
||||
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
const db = await connect(tmpDir.name);
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), false),
|
||||
new Field("text", new Utf8(), false),
|
||||
new Field(
|
||||
"vector",
|
||||
new FixedSizeList(2, new Field("item", new Float32())),
|
||||
false,
|
||||
),
|
||||
]);
|
||||
const data = makeArrowTable(
|
||||
[
|
||||
{ id: 1n, text: "common", vector: [0, 0] },
|
||||
{ id: 2n, text: "common common", vector: [1, 1] },
|
||||
{ id: 3n, text: "common common common", vector: [2, 2] },
|
||||
{ id: 4n, text: "common common common common", vector: [3, 3] },
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
table = await db.createTable("test", data);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
tmpDir.removeCallback();
|
||||
});
|
||||
|
||||
it("applies offset after the vector search limit", async () => {
|
||||
const allResults = await table
|
||||
.vectorSearch([0, 0])
|
||||
.select(["id"])
|
||||
.limit(4)
|
||||
.toArray();
|
||||
const secondPage = await table
|
||||
.vectorSearch([0, 0])
|
||||
.select(["id"])
|
||||
.limit(2)
|
||||
.offset(2)
|
||||
.toArray();
|
||||
|
||||
expect(allResults).toHaveLength(4);
|
||||
expect(secondPage).toHaveLength(2);
|
||||
expect(secondPage.map((row) => row.id)).toEqual(
|
||||
allResults.slice(2, 4).map((row) => row.id),
|
||||
);
|
||||
});
|
||||
|
||||
it("applies offset after the full-text search limit", async () => {
|
||||
await table.createIndex("text", { config: Index.fts() });
|
||||
|
||||
const allResults = await table
|
||||
.search("common", "fts")
|
||||
.select(["id"])
|
||||
.limit(4)
|
||||
.toArray();
|
||||
const secondPage = await table
|
||||
.search("common", "fts")
|
||||
.select(["id"])
|
||||
.limit(2)
|
||||
.offset(2)
|
||||
.toArray();
|
||||
|
||||
expect(allResults).toHaveLength(4);
|
||||
expect(secondPage).toHaveLength(2);
|
||||
expect(secondPage.map((row) => row.id)).toEqual(
|
||||
allResults.slice(2, 4).map((row) => row.id),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Query orderBy", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let table: Table;
|
||||
|
||||
beforeEach(async () => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
const db = await connect(tmpDir.name);
|
||||
|
||||
// Create table with numeric data for sorting
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
new Field("score", new Float32(), true),
|
||||
new Field("name", new Utf8(), true),
|
||||
]);
|
||||
|
||||
const data = makeArrowTable(
|
||||
[
|
||||
{ id: 1n, score: 3.5, name: "charlie" },
|
||||
{ id: 2n, score: 1.2, name: "alice" },
|
||||
{ id: 3n, score: 2.8, name: "bob" },
|
||||
{ id: 4n, score: 0.5, name: "david" },
|
||||
{ id: 5n, score: 4.1, name: "eve" },
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
table = await db.createTable("test", data);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
tmpDir.removeCallback();
|
||||
});
|
||||
|
||||
it("should sort by single column ascending", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.orderBy({ columnName: "score", ascending: true, nullsFirst: false })
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(5);
|
||||
// Verify ascending order
|
||||
expect(results[0].score).toBeCloseTo(0.5, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(1.2, 0.001);
|
||||
expect(results[2].score).toBeCloseTo(2.8, 0.001);
|
||||
expect(results[3].score).toBeCloseTo(3.5, 0.001);
|
||||
expect(results[4].score).toBeCloseTo(4.1, 0.001);
|
||||
});
|
||||
|
||||
it("should sort by single column descending", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.orderBy({ columnName: "score", ascending: false, nullsFirst: false })
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(5);
|
||||
// Verify descending order
|
||||
expect(results[0].score).toBeCloseTo(4.1, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(3.5, 0.001);
|
||||
expect(results[2].score).toBeCloseTo(2.8, 0.001);
|
||||
expect(results[3].score).toBeCloseTo(1.2, 0.001);
|
||||
expect(results[4].score).toBeCloseTo(0.5, 0.001);
|
||||
});
|
||||
|
||||
it("should use ascending as default direction", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.orderBy({ columnName: "score" })
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(5);
|
||||
// Verify ascending order (default)
|
||||
expect(results[0].score).toBeCloseTo(0.5, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(1.2, 0.001);
|
||||
expect(results[2].score).toBeCloseTo(2.8, 0.001);
|
||||
expect(results[3].score).toBeCloseTo(3.5, 0.001);
|
||||
expect(results[4].score).toBeCloseTo(4.1, 0.001);
|
||||
});
|
||||
|
||||
it("should sort by string column", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.orderBy({ columnName: "name" })
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(5);
|
||||
// Verify alphabetical order
|
||||
expect(results[0].name).toBe("alice");
|
||||
expect(results[1].name).toBe("bob");
|
||||
expect(results[2].name).toBe("charlie");
|
||||
expect(results[3].name).toBe("david");
|
||||
expect(results[4].name).toBe("eve");
|
||||
});
|
||||
|
||||
it("should support method chaining with where", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.where("score > 2.0")
|
||||
.orderBy({ columnName: "score" })
|
||||
.toArray();
|
||||
expect(results.length).toBe(3);
|
||||
// Verify filtered and sorted
|
||||
expect(results[0].score).toBeCloseTo(2.8, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(3.5, 0.001);
|
||||
expect(results[2].score).toBeCloseTo(4.1, 0.001);
|
||||
});
|
||||
|
||||
it("should combine repeated where clauses with AND", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.where("score > 1.0")
|
||||
.where("score < 3.0")
|
||||
.orderBy({ columnName: "score" })
|
||||
.toArray();
|
||||
// Only rows matching both predicates should be returned, rather than the
|
||||
// second where() silently replacing the first.
|
||||
expect(results.length).toBe(2);
|
||||
expect(results[0].score).toBeCloseTo(1.2, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(2.8, 0.001);
|
||||
});
|
||||
|
||||
it("should support method chaining with limit", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.orderBy({ columnName: "score", ascending: false })
|
||||
.limit(3)
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(3);
|
||||
// Verify top 3 in descending order
|
||||
expect(results[0].score).toBeCloseTo(4.1, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(3.5, 0.001);
|
||||
expect(results[2].score).toBeCloseTo(2.8, 0.001);
|
||||
});
|
||||
|
||||
it("should support method chaining with offset", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.orderBy({ columnName: "score" })
|
||||
.offset(2)
|
||||
.limit(2)
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(2);
|
||||
// Verify results skip first 2 and take next 2
|
||||
expect(results[0].score).toBeCloseTo(2.8, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(3.5, 0.001);
|
||||
});
|
||||
|
||||
it("should support method chaining with select", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.orderBy({ columnName: "name" })
|
||||
.select(["name", "score"])
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(5);
|
||||
// Verify only selected columns are present
|
||||
expect(Object.keys(results[0])).toEqual(["name", "score"]);
|
||||
expect(Object.keys(results[4])).toEqual(["name", "score"]);
|
||||
// Verify sorted by name
|
||||
expect(results[0].name).toBe("alice");
|
||||
expect(results[4].name).toBe("eve");
|
||||
});
|
||||
|
||||
it("should support complex method chaining", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.where("score > 1.0")
|
||||
.orderBy({ columnName: "score", ascending: false })
|
||||
.limit(3)
|
||||
.select(["id", "score", "name"])
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(3);
|
||||
// Verify filtered, sorted, limited, and projected
|
||||
expect(results[0].score).toBeCloseTo(4.1, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(3.5, 0.001);
|
||||
expect(results[2].score).toBeCloseTo(2.8, 0.001);
|
||||
expect(Object.keys(results[0])).toEqual(["id", "score", "name"]);
|
||||
});
|
||||
|
||||
it("should support multi-column ordering and null placement", async () => {
|
||||
const schema = new Schema([
|
||||
new Field("group", new Int64(), true),
|
||||
new Field("score", new Float32(), true),
|
||||
new Field("name", new Utf8(), true),
|
||||
]);
|
||||
|
||||
const data = makeArrowTable(
|
||||
[
|
||||
{ group: 1n, score: null, name: "z" },
|
||||
{ group: 1n, score: 1.0, name: "b" },
|
||||
{ group: 1n, score: 1.0, name: "a" },
|
||||
{ group: 2n, score: 0.5, name: "c" },
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
const nullTable = await (await connect(tmpDir.name)).createTable(
|
||||
"test_multi_order",
|
||||
data,
|
||||
{ mode: "overwrite" },
|
||||
);
|
||||
|
||||
const results = await nullTable
|
||||
.query()
|
||||
.orderBy([
|
||||
{ columnName: "group", ascending: true, nullsFirst: false },
|
||||
{ columnName: "score", ascending: true, nullsFirst: true },
|
||||
{ columnName: "name", ascending: true, nullsFirst: false },
|
||||
])
|
||||
.toArray();
|
||||
|
||||
expect(results.map((r) => [r.group, r.score, r.name])).toEqual([
|
||||
[1n, null, "z"],
|
||||
[1n, 1.0, "a"],
|
||||
[1n, 1.0, "b"],
|
||||
[2n, 0.5, "c"],
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("IVF search effort", () => {
|
||||
let directory: tmp.DirResult;
|
||||
let table: Table;
|
||||
const vector = Array.from({ length: 16 }, (_, i) => (i === 0 ? 1 : 0));
|
||||
|
||||
beforeEach(async () => {
|
||||
directory = tmp.dirSync({ unsafeCleanup: true });
|
||||
const db = await connect(directory.name);
|
||||
const rows = Array.from({ length: 512 }, (_, id) => ({
|
||||
id,
|
||||
vector: Array.from(
|
||||
{ length: 16 },
|
||||
(_, j) =>
|
||||
(Math.floor(id / 32) === j ? 1 : 0) + Math.sin(id * 17 + j) * 0.001,
|
||||
),
|
||||
}));
|
||||
table = await db.createTable("effort", rows.slice(0, 256));
|
||||
await table.add(rows.slice(256));
|
||||
await table.createIndex("vector", {
|
||||
config: Index.ivfFlat({ numPartitions: 16 }),
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
directory.removeCallback();
|
||||
});
|
||||
|
||||
it.each([
|
||||
[null, 6],
|
||||
[0, 1],
|
||||
[0.25, 3],
|
||||
[0.5, 6],
|
||||
[0.75, 10],
|
||||
[1, 16],
|
||||
])(
|
||||
"preserves effort %s through execution and probe precedence",
|
||||
async (effort, expected) => {
|
||||
const partitions = async (query: ReturnType<Table["vectorSearch"]>) => {
|
||||
const plan = await query.limit(8).analyzePlan();
|
||||
const matches = [...plan.matchAll(/partitions_searched=(\d+)/g)];
|
||||
expect(matches.length).toBeGreaterThan(0);
|
||||
return matches.reduce((total, match) => total + Number(match[1]), 0);
|
||||
};
|
||||
expect(
|
||||
await partitions(table.vectorSearch(vector).searchEffort(effort)),
|
||||
).toBe(expected);
|
||||
expect(
|
||||
await partitions(
|
||||
table.vectorSearch(vector).nprobes(1).searchEffort(effort),
|
||||
),
|
||||
).toBe(effort === null ? 1 : expected);
|
||||
expect(
|
||||
await partitions(
|
||||
table.vectorSearch(vector).searchEffort(effort).nprobes(1),
|
||||
),
|
||||
).toBe(effort === null ? 1 : expected);
|
||||
expect(
|
||||
await partitions(
|
||||
table
|
||||
.vectorSearch(vector)
|
||||
.searchEffort(effort)
|
||||
.nprobes(1)
|
||||
.searchEffort(null),
|
||||
),
|
||||
).toBe(1);
|
||||
expect(
|
||||
await partitions(
|
||||
table.vectorSearch(vector).searchEffort(effort).maximumNprobes(2),
|
||||
),
|
||||
).toBeLessThanOrEqual(2);
|
||||
const actual = await table
|
||||
.vectorSearch(vector)
|
||||
.searchEffort(effort)
|
||||
.limit(8)
|
||||
.toArray();
|
||||
const exact = await table
|
||||
.vectorSearch(vector)
|
||||
.bypassVectorIndex()
|
||||
.limit(8)
|
||||
.toArray();
|
||||
const ids = new Set(exact.map((row) => row.id));
|
||||
expect(
|
||||
actual.filter((row) => ids.has(row.id)).length / 8,
|
||||
).toBeGreaterThanOrEqual(0.5);
|
||||
},
|
||||
);
|
||||
|
||||
it.each([
|
||||
-0.1,
|
||||
1.1,
|
||||
Number.NaN,
|
||||
Number.POSITIVE_INFINITY,
|
||||
Number.NEGATIVE_INFINITY,
|
||||
])("rejects invalid effort %s before execution", (effort) => {
|
||||
expect(() => table.vectorSearch(vector).searchEffort(effort)).toThrow(
|
||||
"search_effort must be finite and in [0, 1]",
|
||||
);
|
||||
expect(() =>
|
||||
table.vectorSearch(vector).fullTextSearch("dog").searchEffort(effort),
|
||||
).toThrow("search_effort must be finite and in [0, 1]");
|
||||
});
|
||||
});
|
||||
@@ -3,13 +3,10 @@
|
||||
|
||||
import * as apiArrow from "apache-arrow";
|
||||
|
||||
import * as arrow15 from "apache-arrow-15";
|
||||
import * as arrow16 from "apache-arrow-16";
|
||||
import * as arrow17 from "apache-arrow-17";
|
||||
import * as arrow18 from "apache-arrow-18";
|
||||
|
||||
import * as tmp from "tmp";
|
||||
|
||||
import { arrowVersions } from "./arrow_versions";
|
||||
|
||||
import { connect } from "../lancedb";
|
||||
import {
|
||||
EmbeddingFunction,
|
||||
@@ -18,7 +15,7 @@ import {
|
||||
} from "../lancedb/embedding";
|
||||
import { getRegistry, register } from "../lancedb/embedding/registry";
|
||||
|
||||
describe.each([arrow15, arrow16, arrow17, arrow18])("LanceSchema", (arrow) => {
|
||||
describe.each(arrowVersions)("LanceSchema", (arrow) => {
|
||||
test("should preserve input order", async () => {
|
||||
const schema = LanceSchema({
|
||||
id: new arrow.Int32(),
|
||||
@@ -29,7 +26,7 @@ describe.each([arrow15, arrow16, arrow17, arrow18])("LanceSchema", (arrow) => {
|
||||
});
|
||||
});
|
||||
|
||||
describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => {
|
||||
describe.each(arrowVersions)("Registry", (arrow) => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
beforeEach(() => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
@@ -106,6 +103,77 @@ describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => {
|
||||
'Embedding function with alias "mock-embedding" already exists',
|
||||
);
|
||||
});
|
||||
test("parseFunctions keeps entries sharing a function name", async () => {
|
||||
class MockEmbeddingFunction extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 3;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow.Float32() as apiArrow.Float;
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return data.map(() => [1, 2, 3]);
|
||||
}
|
||||
}
|
||||
register("mock-embedding")(MockEmbeddingFunction);
|
||||
const parsed = await getRegistry().parseFunctions(
|
||||
new Map([
|
||||
[
|
||||
"embedding_functions",
|
||||
JSON.stringify([
|
||||
{
|
||||
name: "mock-embedding",
|
||||
sourceColumn: "text",
|
||||
vectorColumn: "vector_a",
|
||||
model: {},
|
||||
},
|
||||
{
|
||||
name: "mock-embedding",
|
||||
sourceColumn: "text",
|
||||
vectorColumn: "vector_b",
|
||||
model: {},
|
||||
},
|
||||
]),
|
||||
],
|
||||
]),
|
||||
);
|
||||
expect([...parsed.values()].map((f) => f.vectorColumn)).toEqual([
|
||||
"vector_a",
|
||||
"vector_b",
|
||||
]);
|
||||
|
||||
// The Python bindings write snake_case keys.
|
||||
const snake = await getRegistry().parseFunctions(
|
||||
new Map([
|
||||
[
|
||||
"embedding_functions",
|
||||
JSON.stringify([
|
||||
{
|
||||
name: "mock-embedding",
|
||||
// biome-ignore lint/style/useNamingConvention: the Python wire spelling
|
||||
source_column: "text",
|
||||
// biome-ignore lint/style/useNamingConvention: the Python wire spelling
|
||||
vector_column: "vector_a",
|
||||
model: {},
|
||||
},
|
||||
{
|
||||
name: "mock-embedding",
|
||||
// biome-ignore lint/style/useNamingConvention: the Python wire spelling
|
||||
source_column: "text",
|
||||
// biome-ignore lint/style/useNamingConvention: the Python wire spelling
|
||||
vector_column: "vector_b",
|
||||
model: {},
|
||||
},
|
||||
]),
|
||||
],
|
||||
]),
|
||||
);
|
||||
expect([...snake.keys()]).toEqual(["vector_a", "vector_b"]);
|
||||
expect([...snake.values()].map((f) => f.sourceColumn)).toEqual([
|
||||
"text",
|
||||
"text",
|
||||
]);
|
||||
});
|
||||
test("schema should contain correct metadata", async () => {
|
||||
class MockEmbeddingFunction extends EmbeddingFunction<string> {
|
||||
constructor(args: FunctionOptions = {}) {
|
||||
@@ -169,7 +237,7 @@ describe("Registry.setVar", () => {
|
||||
return data.map(() => [1, 2, 3]);
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow18.Float32() as apiArrow.Float;
|
||||
return new apiArrow.Float32();
|
||||
}
|
||||
protected getSensitiveKeys() {
|
||||
return ["secretKey"];
|
||||
|
||||
+1719
-3
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,8 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { spawn } from "node:child_process";
|
||||
import * as path from "node:path";
|
||||
import { RecordBatch } from "apache-arrow";
|
||||
import * as tmp from "tmp";
|
||||
import { Connection, Index, Table, connect, makeArrowTable } from "../lancedb";
|
||||
@@ -76,4 +78,91 @@ describe("rerankers", function () {
|
||||
|
||||
expect(result).toHaveLength(2);
|
||||
});
|
||||
|
||||
it("does not keep process alive after rerank query", async function () {
|
||||
const script = `
|
||||
import * as lancedb from "./dist/index.js";
|
||||
import * as os from "node:os";
|
||||
import * as path from "node:path";
|
||||
import * as fs from "node:fs/promises";
|
||||
|
||||
const dir = await fs.mkdtemp(path.join(os.tmpdir(), "lancedb-rerank-exit-"));
|
||||
const db = await lancedb.connect(dir);
|
||||
const table = await db.createTable("test", [{ text: "hello", vector: [1, 2, 3] }], {
|
||||
mode: "overwrite",
|
||||
});
|
||||
await table.createIndex("text", { config: lancedb.Index.fts() });
|
||||
await table.waitForIndex(["text_idx"], 30);
|
||||
|
||||
const reranker = await lancedb.rerankers.RRFReranker.create();
|
||||
await table
|
||||
.query()
|
||||
.nearestTo([1, 2, 3])
|
||||
.fullTextSearch("hello")
|
||||
.rerank(reranker)
|
||||
.toArray();
|
||||
|
||||
table.close();
|
||||
db.close();
|
||||
`;
|
||||
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
const child = spawn(
|
||||
process.execPath,
|
||||
["--input-type=module", "-e", script],
|
||||
{
|
||||
cwd: path.resolve(__dirname, ".."),
|
||||
stdio: ["ignore", "pipe", "pipe"],
|
||||
},
|
||||
);
|
||||
|
||||
let stdout = "";
|
||||
let stderr = "";
|
||||
|
||||
child.stdout.on("data", (chunk) => {
|
||||
stdout += chunk.toString();
|
||||
});
|
||||
|
||||
child.stderr.on("data", (chunk) => {
|
||||
stderr += chunk.toString();
|
||||
});
|
||||
|
||||
const timeout = setTimeout(() => {
|
||||
child.kill();
|
||||
reject(
|
||||
new Error(
|
||||
`child process did not exit in time\nstdout:\n${stdout}\nstderr:\n${stderr}`,
|
||||
),
|
||||
);
|
||||
}, 20_000);
|
||||
|
||||
child.on("error", (err) => {
|
||||
clearTimeout(timeout);
|
||||
reject(err);
|
||||
});
|
||||
|
||||
child.on("exit", (code, signal) => {
|
||||
clearTimeout(timeout);
|
||||
if (signal !== null) {
|
||||
reject(
|
||||
new Error(
|
||||
`child process exited with signal ${signal}\nstdout:\n${stdout}\nstderr:\n${stderr}`,
|
||||
),
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if (code !== 0) {
|
||||
reject(
|
||||
new Error(
|
||||
`child process exited with code ${code}\nstdout:\n${stdout}\nstderr:\n${stderr}`,
|
||||
),
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -0,0 +1,307 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as arrow from "../lancedb/arrow";
|
||||
import {
|
||||
sanitizeField,
|
||||
sanitizeMap,
|
||||
sanitizeSchema,
|
||||
sanitizeType,
|
||||
} from "../lancedb/sanitize";
|
||||
import { arrow21 } from "./arrow_versions";
|
||||
|
||||
describe("sanitize", function () {
|
||||
describe("sanitizeType function", function () {
|
||||
it("should handle type objects", function () {
|
||||
const type = new arrow.Int32();
|
||||
const result = sanitizeType(type);
|
||||
|
||||
expect(result.typeId).toBe(arrow.Type.Int);
|
||||
expect((result as arrow.Int).bitWidth).toBe(32);
|
||||
expect((result as arrow.Int).isSigned).toBe(true);
|
||||
|
||||
const floatType = {
|
||||
typeId: 3, // Type.Float = 3
|
||||
precision: 2,
|
||||
toString: () => "Float",
|
||||
isFloat: true,
|
||||
isFixedWidth: true,
|
||||
};
|
||||
|
||||
const floatResult = sanitizeType(floatType);
|
||||
expect(floatResult).toBeInstanceOf(arrow.DataType);
|
||||
expect(floatResult.typeId).toBe(arrow.Type.Float);
|
||||
|
||||
const floatResult2 = sanitizeType({ ...floatType, typeId: () => 3 });
|
||||
expect(floatResult2).toBeInstanceOf(arrow.DataType);
|
||||
expect(floatResult2.typeId).toBe(arrow.Type.Float);
|
||||
});
|
||||
|
||||
const allTypeNameTestCases = [
|
||||
["null", new arrow.Null()],
|
||||
["binary", new arrow.Binary()],
|
||||
["utf8", new arrow.Utf8()],
|
||||
["bool", new arrow.Bool()],
|
||||
["int8", new arrow.Int8()],
|
||||
["int16", new arrow.Int16()],
|
||||
["int32", new arrow.Int32()],
|
||||
["int64", new arrow.Int64()],
|
||||
["uint8", new arrow.Uint8()],
|
||||
["uint16", new arrow.Uint16()],
|
||||
["uint32", new arrow.Uint32()],
|
||||
["uint64", new arrow.Uint64()],
|
||||
["float16", new arrow.Float16()],
|
||||
["float32", new arrow.Float32()],
|
||||
["float64", new arrow.Float64()],
|
||||
["datemillisecond", new arrow.DateMillisecond()],
|
||||
["dateday", new arrow.DateDay()],
|
||||
["timenanosecond", new arrow.TimeNanosecond()],
|
||||
["timemicrosecond", new arrow.TimeMicrosecond()],
|
||||
["timemillisecond", new arrow.TimeMillisecond()],
|
||||
["timesecond", new arrow.TimeSecond()],
|
||||
["intervaldaytime", new arrow.IntervalDayTime()],
|
||||
["intervalyearmonth", new arrow.IntervalYearMonth()],
|
||||
["durationnanosecond", new arrow.DurationNanosecond()],
|
||||
["durationmicrosecond", new arrow.DurationMicrosecond()],
|
||||
["durationmillisecond", new arrow.DurationMillisecond()],
|
||||
["durationsecond", new arrow.DurationSecond()],
|
||||
["largebinary", new arrow.LargeBinary()],
|
||||
["largeutf8", new arrow.LargeUtf8()],
|
||||
] as const;
|
||||
|
||||
it.each(allTypeNameTestCases)(
|
||||
'should map type name "%s" to %s',
|
||||
function (name, expected) {
|
||||
const result = sanitizeType(name);
|
||||
expect(result).toBeInstanceOf(expected.constructor);
|
||||
},
|
||||
);
|
||||
|
||||
const caseVariationTestCases = [
|
||||
["NULL", new arrow.Null()],
|
||||
["Utf8", new arrow.Utf8()],
|
||||
["FLOAT32", new arrow.Float32()],
|
||||
["DaTedAy", new arrow.DateDay()],
|
||||
] as const;
|
||||
|
||||
it.each(caseVariationTestCases)(
|
||||
'should be case insensitive for type name "%s" mapped to %s',
|
||||
function (name, expected) {
|
||||
const result = sanitizeType(name);
|
||||
expect(result).toBeInstanceOf(expected.constructor);
|
||||
},
|
||||
);
|
||||
|
||||
it("should throw error for unrecognized type name", function () {
|
||||
expect(() => sanitizeType("invalid_type")).toThrow(
|
||||
"Unrecognized type name in schema: invalid_type",
|
||||
);
|
||||
});
|
||||
|
||||
// These types only exist in Arrow 20 (IntervalMonthDayNano) and 21 (the
|
||||
// rest), so the objects come from a second copy of the newest release. The
|
||||
// library itself may be running on an older Arrow (see the CI matrix), in
|
||||
// which case sanitizing them has to fail with a clear message instead.
|
||||
const hostArrow = arrow as unknown as Record<string, unknown>;
|
||||
const hostSupportsNewerTypes = typeof hostArrow.Utf8View === "function";
|
||||
const whenSupported = hostSupportsNewerTypes ? describe : describe.skip;
|
||||
const whenUnsupported = hostSupportsNewerTypes ? describe.skip : describe;
|
||||
|
||||
const newerTypeTestCases = [
|
||||
["Utf8View", () => new arrow21.Utf8View()],
|
||||
["BinaryView", () => new arrow21.BinaryView()],
|
||||
["IntervalMonthDayNano", () => new arrow21.IntervalMonthDayNano()],
|
||||
[
|
||||
"LargeList",
|
||||
() =>
|
||||
new arrow21.LargeList(new arrow21.Field("item", new arrow21.Utf8())),
|
||||
],
|
||||
] as const;
|
||||
|
||||
whenSupported("types added in Arrow 20 and 21", function () {
|
||||
it.each(newerTypeTestCases)(
|
||||
"should sanitize a %s type from another copy of Arrow",
|
||||
function (name, makeForeign) {
|
||||
const foreign = makeForeign();
|
||||
const result = sanitizeType(foreign);
|
||||
// Interval subclasses share `Type.Interval`, so the sanitized type
|
||||
// is the generic `Interval` carrying the same unit, as for the
|
||||
// other units.
|
||||
const expected = name === "IntervalMonthDayNano" ? "Interval" : name;
|
||||
expect(result).toBeInstanceOf(hostArrow[expected] as () => unknown);
|
||||
expect(result.typeId).toBe(foreign.typeId);
|
||||
expect(result.toString()).toBe(foreign.toString());
|
||||
},
|
||||
);
|
||||
|
||||
it.each([
|
||||
["binaryview", "BinaryView"],
|
||||
["utf8view", "Utf8View"],
|
||||
["intervalmonthdaynano", "Interval"],
|
||||
])('should map type name "%s" to %s', function (name, expected) {
|
||||
expect(sanitizeType(name)).toBeInstanceOf(
|
||||
hostArrow[expected] as () => unknown,
|
||||
);
|
||||
});
|
||||
|
||||
it("should sanitize the child field of a LargeList", function () {
|
||||
const foreign = new arrow21.LargeList(
|
||||
new arrow21.Field("item", new arrow21.Int32(), false),
|
||||
);
|
||||
const result = sanitizeType(foreign) as arrow.List;
|
||||
expect(result.children[0]).toBeInstanceOf(arrow.Field);
|
||||
expect(result.children[0].name).toBe("item");
|
||||
expect(result.children[0].nullable).toBe(false);
|
||||
const child = result.children[0].type;
|
||||
expect(arrow.DataType.isInt(child)).toBe(true);
|
||||
expect((child as arrow.Int).bitWidth).toBe(32);
|
||||
expect((child as arrow.Int).isSigned).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
whenUnsupported("types newer than the installed Arrow", function () {
|
||||
it.each(
|
||||
newerTypeTestCases.filter(([name]) => name !== "IntervalMonthDayNano"),
|
||||
)(
|
||||
"should explain that %s needs a newer apache-arrow",
|
||||
function (name, makeForeign) {
|
||||
expect(() => sanitizeType(makeForeign())).toThrow(
|
||||
`The Arrow type ${name} requires apache-arrow`,
|
||||
);
|
||||
},
|
||||
);
|
||||
|
||||
it("should sanitize IntervalMonthDayNano to a generic Interval", () => {
|
||||
// The subclass shares `Type.Interval`, which every supported Arrow
|
||||
// release has; only the unit is new.
|
||||
const result = sanitizeType(new arrow21.IntervalMonthDayNano());
|
||||
expect(result).toBeInstanceOf(arrow.Interval);
|
||||
expect((result as arrow.Interval).unit).toBe(
|
||||
arrow21.IntervalUnit.MONTH_DAY_NANO,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it("should rebuild a schema from a second copy of the same Arrow version", function () {
|
||||
// Arrow 21 makes `instanceof Schema` succeed across copies of the library
|
||||
// but not `instanceof <DataType>`; the schema must still be rebuilt.
|
||||
const foreign = new arrow21.Schema([
|
||||
new arrow21.Field("x", new arrow21.Float32(), true),
|
||||
]);
|
||||
// Only Arrow 21+ overrides `Symbol.hasInstance` this way.
|
||||
const crossCopyInstanceof = Object.getOwnPropertySymbols(
|
||||
arrow.Schema,
|
||||
).includes(Symbol.hasInstance);
|
||||
expect(foreign instanceof arrow.Schema).toBe(crossCopyInstanceof);
|
||||
const result = sanitizeSchema(foreign);
|
||||
expect(result).not.toBe(foreign);
|
||||
const type = result.fields[0].type;
|
||||
expect(type).toBeInstanceOf(arrow.DataType);
|
||||
expect(arrow.DataType.isFloat(type)).toBe(true);
|
||||
expect((type as arrow.Float).precision).toBe(arrow.Precision.SINGLE);
|
||||
expect(() => arrow.makeBuilder({ type })).not.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe("sanitizeField function", function () {
|
||||
it("should handle field with string type name", function () {
|
||||
const field = sanitizeField({
|
||||
name: "string_field",
|
||||
type: "utf8",
|
||||
nullable: true,
|
||||
metadata: new Map([["key", "value"]]),
|
||||
});
|
||||
|
||||
expect(field).toBeInstanceOf(arrow.Field);
|
||||
expect(field.name).toBe("string_field");
|
||||
expect(field.type).toBeInstanceOf(arrow.Utf8);
|
||||
expect(field.nullable).toBe(true);
|
||||
expect(field.metadata?.get("key")).toBe("value");
|
||||
});
|
||||
|
||||
it("should handle field with type object", function () {
|
||||
const floatType = {
|
||||
typeId: 3, // Float
|
||||
precision: 32,
|
||||
};
|
||||
|
||||
const field = sanitizeField({
|
||||
name: "float_field",
|
||||
type: floatType,
|
||||
nullable: false,
|
||||
});
|
||||
|
||||
expect(field).toBeInstanceOf(arrow.Field);
|
||||
expect(field.name).toBe("float_field");
|
||||
expect(field.type).toBeInstanceOf(arrow.DataType);
|
||||
expect(field.type.typeId).toBe(arrow.Type.Float);
|
||||
expect((field.type as arrow.Float64).precision).toBe(32);
|
||||
expect(field.nullable).toBe(false);
|
||||
});
|
||||
|
||||
it("should handle field with direct Type instance", function () {
|
||||
const field = sanitizeField({
|
||||
name: "bool_field",
|
||||
type: new arrow.Bool(),
|
||||
nullable: true,
|
||||
});
|
||||
|
||||
expect(field).toBeInstanceOf(arrow.Field);
|
||||
expect(field.name).toBe("bool_field");
|
||||
expect(field.type).toBeInstanceOf(arrow.Bool);
|
||||
expect(field.nullable).toBe(true);
|
||||
});
|
||||
|
||||
it("should throw error for invalid field object", function () {
|
||||
expect(() =>
|
||||
sanitizeField({
|
||||
type: "int32",
|
||||
nullable: true,
|
||||
}),
|
||||
).toThrow(
|
||||
"The field passed in is missing a `type`/`name`/`nullable` property",
|
||||
);
|
||||
|
||||
// Invalid type
|
||||
expect(() =>
|
||||
sanitizeField({
|
||||
name: "invalid",
|
||||
type: { invalid: true },
|
||||
nullable: true,
|
||||
}),
|
||||
).toThrow("Expected a Type to have a typeId property");
|
||||
|
||||
// Invalid nullable
|
||||
expect(() =>
|
||||
sanitizeField({
|
||||
name: "invalid_nullable",
|
||||
type: "int32",
|
||||
nullable: "not a boolean",
|
||||
}),
|
||||
).toThrow("The field passed in had a non-boolean `nullable` property");
|
||||
});
|
||||
|
||||
it("should report error for invalid type name", function () {
|
||||
expect(() =>
|
||||
sanitizeField({
|
||||
name: "invalid_field",
|
||||
type: "invalid_type",
|
||||
nullable: true,
|
||||
}),
|
||||
).toThrow(
|
||||
"Unable to sanitize type for field: invalid_field due to error: Error: Unrecognized type name in schema: invalid_type",
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe("sanitizeMap function", function () {
|
||||
it.each([
|
||||
["no children", []],
|
||||
["two children", [{}, {}]],
|
||||
])("should reject a Map type with %s", function (_, children) {
|
||||
expect(() => sanitizeMap({ children, keysSorted: false })).toThrow(
|
||||
"Expected a Map type to have exactly one child",
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,438 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import {
|
||||
Field,
|
||||
Float16,
|
||||
Int32,
|
||||
type RecordBatch,
|
||||
RecordBatchReader,
|
||||
Schema,
|
||||
tableToIPC,
|
||||
} from "apache-arrow";
|
||||
import { makeArrowTable, makeEmptyTable } from "../lancedb/arrow";
|
||||
import { Scannable } from "../lancedb/scannable";
|
||||
|
||||
function makeTable() {
|
||||
return makeArrowTable(
|
||||
[
|
||||
{ id: 1, name: "a" },
|
||||
{ id: 2, name: "b" },
|
||||
{ id: 3, name: "c" },
|
||||
],
|
||||
{ vectorColumns: {} },
|
||||
);
|
||||
}
|
||||
|
||||
async function makeReader(): Promise<RecordBatchReader> {
|
||||
// `RecordBatchReader.from()` returns an unopened reader; `.schema` is only
|
||||
// populated after `.open()`. Opening sync readers is synchronous.
|
||||
const reader = RecordBatchReader.from(tableToIPC(makeTable()));
|
||||
return reader.open() as RecordBatchReader;
|
||||
}
|
||||
|
||||
describe("Scannable", () => {
|
||||
describe("fromTable", () => {
|
||||
test("reflects schema, numRows, and defaults rescannable=true", async () => {
|
||||
const table = makeTable();
|
||||
const scannable = await Scannable.fromTable(table);
|
||||
|
||||
expect(scannable.schema).toBe(table.schema);
|
||||
expect(scannable.numRows).toBe(table.numRows);
|
||||
expect(scannable.rescannable).toBe(true);
|
||||
});
|
||||
|
||||
test("throws when opts.numRows does not match table.numRows", async () => {
|
||||
await expect(
|
||||
Scannable.fromTable(makeTable(), { numRows: 42 }),
|
||||
).rejects.toThrow(/does not match table\.numRows/);
|
||||
});
|
||||
|
||||
test("throws when opts.rescannable is false", async () => {
|
||||
await expect(
|
||||
Scannable.fromTable(makeTable(), { rescannable: false }),
|
||||
).rejects.toThrow(/always rescannable/);
|
||||
});
|
||||
});
|
||||
|
||||
describe("fromRecordBatchReader", () => {
|
||||
test("reflects schema and defaults numRows=null, rescannable=false", async () => {
|
||||
const reader = await makeReader();
|
||||
const scannable = await Scannable.fromRecordBatchReader(reader);
|
||||
|
||||
expect(scannable.schema).toBe(reader.schema);
|
||||
expect(scannable.numRows).toBeNull();
|
||||
expect(scannable.rescannable).toBe(false);
|
||||
});
|
||||
|
||||
test("honors numRows override", async () => {
|
||||
const scannable = await Scannable.fromRecordBatchReader(
|
||||
await makeReader(),
|
||||
{ numRows: 3 },
|
||||
);
|
||||
|
||||
expect(scannable.numRows).toBe(3);
|
||||
expect(scannable.rescannable).toBe(false);
|
||||
});
|
||||
|
||||
test("rescannable: false explicit does not throw", async () => {
|
||||
const reader = await makeReader();
|
||||
const scannable = await Scannable.fromRecordBatchReader(reader, {
|
||||
rescannable: false,
|
||||
});
|
||||
expect(scannable.rescannable).toBe(false);
|
||||
});
|
||||
|
||||
test("throws when opts.rescannable is true", async () => {
|
||||
const reader = await makeReader();
|
||||
await expect(
|
||||
Scannable.fromRecordBatchReader(reader, { rescannable: true }),
|
||||
).rejects.toThrow(/does not accept rescannable/);
|
||||
});
|
||||
|
||||
test("throws when opts.rescannable is true even alongside numRows", async () => {
|
||||
const reader = await makeReader();
|
||||
await expect(
|
||||
Scannable.fromRecordBatchReader(reader, {
|
||||
numRows: 3,
|
||||
rescannable: true,
|
||||
}),
|
||||
).rejects.toThrow(/does not accept rescannable/);
|
||||
});
|
||||
});
|
||||
|
||||
describe("fromIterable", () => {
|
||||
test("accepts a sync iterable of batches", async () => {
|
||||
const table = makeTable();
|
||||
const scannable = await Scannable.fromIterable(
|
||||
table.schema,
|
||||
table.batches,
|
||||
);
|
||||
|
||||
expect(scannable.schema).toBe(table.schema);
|
||||
expect(scannable.numRows).toBeNull();
|
||||
expect(scannable.rescannable).toBe(false);
|
||||
});
|
||||
|
||||
test("accepts an async iterable of batches", async () => {
|
||||
const table = makeTable();
|
||||
async function* generator(): AsyncGenerator<RecordBatch> {
|
||||
for (const batch of table.batches) {
|
||||
yield batch;
|
||||
}
|
||||
}
|
||||
|
||||
const scannable = await Scannable.fromIterable(table.schema, generator());
|
||||
expect(scannable.schema).toBe(table.schema);
|
||||
expect(scannable.rescannable).toBe(false);
|
||||
});
|
||||
|
||||
describe("rescannable: true detection", () => {
|
||||
// Replayable inputs: [Symbol.iterator]() / [Symbol.asyncIterator]()
|
||||
// returns a fresh iterator each call. Must NOT throw.
|
||||
|
||||
test("Array passes (fresh ArrayIterator each call)", async () => {
|
||||
const table = makeTable();
|
||||
const scannable = await Scannable.fromIterable(
|
||||
table.schema,
|
||||
table.batches,
|
||||
{ rescannable: true },
|
||||
);
|
||||
expect(scannable.rescannable).toBe(true);
|
||||
});
|
||||
|
||||
test("Set passes (fresh SetIterator each call)", async () => {
|
||||
const table = makeTable();
|
||||
const set = new Set<RecordBatch>(table.batches);
|
||||
const scannable = await Scannable.fromIterable(table.schema, set, {
|
||||
rescannable: true,
|
||||
});
|
||||
expect(scannable.rescannable).toBe(true);
|
||||
});
|
||||
|
||||
test("custom Iterable returning a fresh iterator passes", async () => {
|
||||
const table = makeTable();
|
||||
const replayable: Iterable<RecordBatch> = {
|
||||
[Symbol.iterator]() {
|
||||
return table.batches[Symbol.iterator]();
|
||||
},
|
||||
};
|
||||
const scannable = await Scannable.fromIterable(
|
||||
table.schema,
|
||||
replayable,
|
||||
{ rescannable: true },
|
||||
);
|
||||
expect(scannable.rescannable).toBe(true);
|
||||
});
|
||||
|
||||
test("object with generator method passes (fresh generator each call)", async () => {
|
||||
const table = makeTable();
|
||||
const replayable: Iterable<RecordBatch> = {
|
||||
*[Symbol.iterator]() {
|
||||
for (const batch of table.batches) yield batch;
|
||||
},
|
||||
};
|
||||
const scannable = await Scannable.fromIterable(
|
||||
table.schema,
|
||||
replayable,
|
||||
{ rescannable: true },
|
||||
);
|
||||
expect(scannable.rescannable).toBe(true);
|
||||
});
|
||||
|
||||
test("empty Array passes (replayable degenerate case)", async () => {
|
||||
const schema = makeTable().schema;
|
||||
const scannable = await Scannable.fromIterable(
|
||||
schema,
|
||||
[] as RecordBatch[],
|
||||
{ rescannable: true },
|
||||
);
|
||||
expect(scannable.rescannable).toBe(true);
|
||||
});
|
||||
|
||||
// One-shot inputs: [Symbol.iterator]() / [Symbol.asyncIterator]()
|
||||
// returns the same object, or the input is already-an-iterator.
|
||||
// Must throw with a /one-shot/ message.
|
||||
|
||||
test("sync generator throws", async () => {
|
||||
const table = makeTable();
|
||||
function* generator(): Generator<RecordBatch> {
|
||||
for (const batch of table.batches) yield batch;
|
||||
}
|
||||
await expect(
|
||||
Scannable.fromIterable(table.schema, generator(), {
|
||||
rescannable: true,
|
||||
}),
|
||||
).rejects.toThrow(/one-shot/);
|
||||
});
|
||||
|
||||
test("async generator throws", async () => {
|
||||
const table = makeTable();
|
||||
async function* generator(): AsyncGenerator<RecordBatch> {
|
||||
for (const batch of table.batches) yield batch;
|
||||
}
|
||||
await expect(
|
||||
Scannable.fromIterable(table.schema, generator(), {
|
||||
rescannable: true,
|
||||
}),
|
||||
).rejects.toThrow(/one-shot/);
|
||||
});
|
||||
|
||||
test("empty generator throws (one-shot degenerate case)", async () => {
|
||||
const schema = makeTable().schema;
|
||||
function* generator(): Generator<RecordBatch> {
|
||||
// intentionally empty; yields nothing.
|
||||
}
|
||||
await expect(
|
||||
Scannable.fromIterable(schema, generator(), { rescannable: true }),
|
||||
).rejects.toThrow(/one-shot/);
|
||||
});
|
||||
|
||||
test("custom self-iterator throws", async () => {
|
||||
const table = makeTable();
|
||||
const batches = table.batches;
|
||||
let i = 0;
|
||||
const oneShot: Iterable<RecordBatch> & Iterator<RecordBatch> = {
|
||||
[Symbol.iterator]() {
|
||||
return this;
|
||||
},
|
||||
next() {
|
||||
if (i >= batches.length) {
|
||||
return { done: true, value: undefined };
|
||||
}
|
||||
return { done: false, value: batches[i++] };
|
||||
},
|
||||
};
|
||||
await expect(
|
||||
Scannable.fromIterable(table.schema, oneShot, { rescannable: true }),
|
||||
).rejects.toThrow(/one-shot/);
|
||||
});
|
||||
|
||||
test("Array.values() (IterableIterator) throws", async () => {
|
||||
const table = makeTable();
|
||||
const iter = table.batches.values();
|
||||
await expect(
|
||||
Scannable.fromIterable(table.schema, iter, { rescannable: true }),
|
||||
).rejects.toThrow(/one-shot/);
|
||||
});
|
||||
|
||||
test("raw iterator (only `.next`) throws", async () => {
|
||||
const table = makeTable();
|
||||
const batches = table.batches;
|
||||
let i = 0;
|
||||
const rawIter = {
|
||||
next(): IteratorResult<RecordBatch> {
|
||||
if (i >= batches.length) {
|
||||
return { done: true, value: undefined };
|
||||
}
|
||||
return { done: false, value: batches[i++] };
|
||||
},
|
||||
};
|
||||
await expect(
|
||||
Scannable.fromIterable(
|
||||
table.schema,
|
||||
rawIter as unknown as Iterable<RecordBatch>,
|
||||
{ rescannable: true },
|
||||
),
|
||||
).rejects.toThrow(/one-shot/);
|
||||
});
|
||||
|
||||
// Edge: null/undefined must not crash the detection helper. The
|
||||
// null check belongs to `normalizeIterator` and only fires when a
|
||||
// scan starts.
|
||||
|
||||
test("null input does not crash detection at construction", async () => {
|
||||
const schema = makeTable().schema;
|
||||
await expect(
|
||||
Scannable.fromIterable(
|
||||
schema,
|
||||
null as unknown as Iterable<RecordBatch>,
|
||||
{
|
||||
rescannable: true,
|
||||
},
|
||||
),
|
||||
).resolves.toBeDefined();
|
||||
});
|
||||
|
||||
test("undefined input does not crash detection at construction", async () => {
|
||||
const schema = makeTable().schema;
|
||||
await expect(
|
||||
Scannable.fromIterable(
|
||||
schema,
|
||||
undefined as unknown as Iterable<RecordBatch>,
|
||||
{ rescannable: true },
|
||||
),
|
||||
).resolves.toBeDefined();
|
||||
});
|
||||
|
||||
// Default (rescannable omitted) skips the check entirely, so even
|
||||
// pathological inputs construct without throwing here.
|
||||
|
||||
test("rescannable omitted skips detection entirely (generator passes)", async () => {
|
||||
const table = makeTable();
|
||||
function* generator(): Generator<RecordBatch> {
|
||||
for (const batch of table.batches) yield batch;
|
||||
}
|
||||
const scannable = await Scannable.fromIterable(
|
||||
table.schema,
|
||||
generator(),
|
||||
);
|
||||
expect(scannable.rescannable).toBe(false);
|
||||
});
|
||||
|
||||
test("rescannable: false explicit skips detection entirely (generator passes)", async () => {
|
||||
const table = makeTable();
|
||||
function* generator(): Generator<RecordBatch> {
|
||||
for (const batch of table.batches) yield batch;
|
||||
}
|
||||
const scannable = await Scannable.fromIterable(
|
||||
table.schema,
|
||||
generator(),
|
||||
{ rescannable: false },
|
||||
);
|
||||
expect(scannable.rescannable).toBe(false);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("fromFactory", () => {
|
||||
test("defaults rescannable=true and does not invoke the factory eagerly", async () => {
|
||||
const table = makeTable();
|
||||
const factory = jest.fn(() => table.batches);
|
||||
|
||||
const scannable = await Scannable.fromFactory(table.schema, factory);
|
||||
|
||||
expect(scannable.schema).toBe(table.schema);
|
||||
expect(scannable.rescannable).toBe(true);
|
||||
expect(factory).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test("honors rescannable and numRows overrides", async () => {
|
||||
const table = makeTable();
|
||||
const scannable = await Scannable.fromFactory(
|
||||
table.schema,
|
||||
() => table.batches,
|
||||
{ numRows: 7, rescannable: false },
|
||||
);
|
||||
|
||||
expect(scannable.numRows).toBe(7);
|
||||
expect(scannable.rescannable).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("validation", () => {
|
||||
test("throws when numRows is negative", async () => {
|
||||
await expect(
|
||||
Scannable.fromFactory(makeTable().schema, () => [], { numRows: -1 }),
|
||||
).rejects.toThrow(/non-negative/);
|
||||
});
|
||||
|
||||
test("throws when numRows is not an integer", async () => {
|
||||
await expect(
|
||||
Scannable.fromFactory(makeTable().schema, () => [], { numRows: 3.5 }),
|
||||
).rejects.toThrow(/integer/);
|
||||
});
|
||||
});
|
||||
|
||||
describe("native handle", () => {
|
||||
test("exposes a native handle via inner", async () => {
|
||||
const scannable = await Scannable.fromTable(makeTable());
|
||||
expect(scannable.inner).toBeDefined();
|
||||
expect(typeof scannable.inner).toBe("object");
|
||||
expect(scannable.inner).not.toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
// Schema-variety construction tests. Each asserts that construction
|
||||
// succeeds against a richer Arrow schema, which transitively exercises
|
||||
// schema serialization and the Rust-side `ipc_file_to_schema` for types
|
||||
// beyond flat primitives.
|
||||
describe("schema variety", () => {
|
||||
test("accepts an empty table", async () => {
|
||||
const schema = new Schema([new Field("id", new Int32(), true)]);
|
||||
const table = makeEmptyTable(schema);
|
||||
const scannable = await Scannable.fromTable(table);
|
||||
|
||||
expect(scannable.numRows).toBe(0);
|
||||
expect(scannable.schema).toBe(table.schema);
|
||||
});
|
||||
|
||||
test("accepts nested struct and list columns", async () => {
|
||||
const table = makeArrowTable(
|
||||
[
|
||||
{ id: 1, point: { x: 0, y: 0 }, tags: ["a", "b"] },
|
||||
{ id: 2, point: { x: 1, y: 2 }, tags: ["c"] },
|
||||
],
|
||||
{ vectorColumns: {} },
|
||||
);
|
||||
const scannable = await Scannable.fromTable(table);
|
||||
|
||||
expect(scannable.schema).toBe(table.schema);
|
||||
expect(scannable.numRows).toBe(2);
|
||||
});
|
||||
|
||||
test("accepts a FixedSizeList (vector) column", async () => {
|
||||
const table = makeArrowTable(
|
||||
[
|
||||
{ id: 1, vec: [1, 2, 3] },
|
||||
{ id: 2, vec: [4, 5, 6] },
|
||||
],
|
||||
{ vectorColumns: { vec: { type: new Float16() } } },
|
||||
);
|
||||
const scannable = await Scannable.fromTable(table);
|
||||
|
||||
expect(scannable.schema).toBe(table.schema);
|
||||
expect(scannable.numRows).toBe(2);
|
||||
});
|
||||
|
||||
test("accepts a table with many columns", async () => {
|
||||
const row: Record<string, number> = {};
|
||||
for (let i = 0; i < 50; i++) row[`c${i}`] = i;
|
||||
const table = makeArrowTable([row, row], { vectorColumns: {} });
|
||||
const scannable = await Scannable.fromTable(table);
|
||||
|
||||
expect(scannable.schema.fields.length).toBe(50);
|
||||
expect(scannable.numRows).toBe(2);
|
||||
});
|
||||
});
|
||||
});
|
||||
+3478
-45
File diff suppressed because it is too large
Load Diff
@@ -14,6 +14,9 @@ test.each([
|
||||
[new ArrayBuffer(8), "X'0000000000000000'"],
|
||||
[Buffer.from("hello"), "X'68656c6c6f'"],
|
||||
["Hello 'world'", "'Hello ''world'''"],
|
||||
[NaN, "CAST('NaN' AS DOUBLE)"],
|
||||
[Infinity, "CAST('Infinity' AS DOUBLE)"],
|
||||
[-Infinity, "CAST('-Infinity' AS DOUBLE)"],
|
||||
])("toSQL(%p) === %p", (value, expected) => {
|
||||
expect(toSQL(value)).toBe(expected);
|
||||
});
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import * as tmp from "tmp";
|
||||
|
||||
import { type Table, connect } from "../lancedb";
|
||||
import {
|
||||
Field,
|
||||
FixedSizeList,
|
||||
Float32,
|
||||
Int64,
|
||||
Schema,
|
||||
makeArrowTable,
|
||||
} from "../lancedb/arrow";
|
||||
|
||||
describe("Vector query with different typed arrays", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
|
||||
afterEach(() => {
|
||||
tmpDir?.removeCallback();
|
||||
});
|
||||
|
||||
async function createFloat32Table(): Promise<Table> {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
const db = await connect(tmpDir.name);
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
new Field(
|
||||
"vec",
|
||||
new FixedSizeList(2, new Field("item", new Float32())),
|
||||
true,
|
||||
),
|
||||
]);
|
||||
const data = makeArrowTable(
|
||||
[
|
||||
{ id: 1n, vec: [1.0, 0.0] },
|
||||
{ id: 2n, vec: [0.0, 1.0] },
|
||||
{ id: 3n, vec: [1.0, 1.0] },
|
||||
],
|
||||
{ schema },
|
||||
);
|
||||
return db.createTable("test_f32", data);
|
||||
}
|
||||
|
||||
it("should search with Float32Array (baseline)", async () => {
|
||||
const table = await createFloat32Table();
|
||||
const results = await table
|
||||
.query()
|
||||
.nearestTo(new Float32Array([1.0, 0.0]))
|
||||
.limit(1)
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(1);
|
||||
expect(Number(results[0].id)).toBe(1);
|
||||
});
|
||||
|
||||
it("should search with number[] (backward compat)", async () => {
|
||||
const table = await createFloat32Table();
|
||||
const results = await table
|
||||
.query()
|
||||
.nearestTo([1.0, 0.0])
|
||||
.limit(1)
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(1);
|
||||
expect(Number(results[0].id)).toBe(1);
|
||||
});
|
||||
|
||||
it("should search with Float64Array via raw path", async () => {
|
||||
const table = await createFloat32Table();
|
||||
const results = await table
|
||||
.query()
|
||||
.nearestTo(new Float64Array([1.0, 0.0]))
|
||||
.limit(1)
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(1);
|
||||
expect(Number(results[0].id)).toBe(1);
|
||||
});
|
||||
|
||||
it("should add multiple query vectors with Float64Array", async () => {
|
||||
const table = await createFloat32Table();
|
||||
const results = await table
|
||||
.query()
|
||||
.nearestTo(new Float64Array([1.0, 0.0]))
|
||||
.addQueryVector(new Float64Array([0.0, 1.0]))
|
||||
.limit(2)
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBeGreaterThanOrEqual(2);
|
||||
});
|
||||
|
||||
// Float16Array is only available in Node 22+; not in TypeScript's standard lib yet
|
||||
const float16ArrayCtor = (globalThis as unknown as Record<string, unknown>)
|
||||
.Float16Array as (new (values: number[]) => unknown) | undefined;
|
||||
const hasFloat16 = float16ArrayCtor !== undefined;
|
||||
const f16it = hasFloat16 ? it : it.skip;
|
||||
|
||||
f16it("should search with Float16Array via raw path", async () => {
|
||||
const table = await createFloat32Table();
|
||||
const results = await table
|
||||
.query()
|
||||
.nearestTo(new float16ArrayCtor!([1.0, 0.0]) as Float32Array)
|
||||
.limit(1)
|
||||
.toArray();
|
||||
|
||||
expect(results.length).toBe(1);
|
||||
expect(Number(results[0].id)).toBe(1);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user