Files
lancedb/nodejs/__test__/blob.test.ts
T

696 lines
22 KiB
TypeScript

// 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");
});
});