mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-20 18:52:28 +00:00
fix!: normalize blob layouts and support external byte ranges (#4589)
This commit is contained in:
@@ -11,6 +11,15 @@ function blob(name, options): Field
|
||||
```
|
||||
|
||||
Declares a `lance.blob.v2` column.
|
||||
Uses nullable `data`, `uri`, `position`, and `size` fields for inline bytes,
|
||||
whole external files, and external byte ranges. Set both `position` and
|
||||
positive `size` for a range, or omit both for a whole file.
|
||||
|
||||
Existing two-field columns remain writable. Appends and merge inserts fill
|
||||
missing range fields with nulls for four-field targets and drop all-null
|
||||
range fields for two-field targets. Non-null ranges require a four-field
|
||||
target. Overwrite with Arrow input replaces the table's schema; row objects
|
||||
are converted using the current table schema.
|
||||
|
||||
Query results are descriptors, not payload bytes. Use [Table.fetchBlobs](../classes/Table.md#fetchblobs)
|
||||
or [Table.fetchBlobFiles](../classes/Table.md#fetchblobfiles) to read bytes.
|
||||
|
||||
@@ -12,6 +12,7 @@ type BlobInput:
|
||||
| BlobUri
|
||||
| object
|
||||
| object
|
||||
| object
|
||||
| null
|
||||
| undefined;
|
||||
```
|
||||
@@ -20,6 +21,9 @@ A value accepted for a `lance.blob.v2` column (see [blob](../functions/blob.md))
|
||||
|
||||
Either inline bytes, a URI pointing at external bytes, or a struct that sets
|
||||
exactly one of `data` or `uri`. `null` and `undefined` write a null blob.
|
||||
External byte ranges require both
|
||||
`position` (non-negative) and `size` (positive). Use safe integer numbers or
|
||||
`bigint` values for these unsigned 64-bit byte counts.
|
||||
|
||||
## Example
|
||||
|
||||
|
||||
@@ -21,7 +21,41 @@ 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);
|
||||
@@ -156,7 +190,7 @@ describe("coerceBlobValue", () => {
|
||||
[
|
||||
"unknown struct field",
|
||||
{ data: Buffer.from("y"), mime: "image/png" },
|
||||
/Blob struct values only support 'data' and 'uri', got 'mime'/,
|
||||
/Blob struct values only support .*got 'mime'/,
|
||||
],
|
||||
[
|
||||
"misspelled struct field",
|
||||
@@ -174,6 +208,158 @@ describe("coerceBlobValue", () => {
|
||||
});
|
||||
});
|
||||
|
||||
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([
|
||||
|
||||
@@ -12,6 +12,7 @@ import {
|
||||
OAuthConfig,
|
||||
OAuthFlowType,
|
||||
TlsConfig,
|
||||
blob,
|
||||
connect,
|
||||
} from "../lancedb";
|
||||
import {
|
||||
@@ -19,7 +20,10 @@ import {
|
||||
Field,
|
||||
Int64,
|
||||
RecordBatch,
|
||||
RecordBatchReader,
|
||||
Schema,
|
||||
Struct,
|
||||
makeArrowTable,
|
||||
tableToIPC,
|
||||
} from "../lancedb/arrow";
|
||||
import {
|
||||
@@ -94,6 +98,137 @@ async function withMockDatabase(
|
||||
}
|
||||
|
||||
describe("remote connection", () => {
|
||||
it.each([false, true])(
|
||||
"normalizes remote blob writes (range target: %s)",
|
||||
async (wide) => {
|
||||
const field = (full: boolean) => {
|
||||
const f = blob("image");
|
||||
return full
|
||||
? f
|
||||
: new Field(
|
||||
f.name,
|
||||
new Struct((f.type as Struct).children.slice(0, 2)),
|
||||
f.nullable,
|
||||
f.metadata,
|
||||
);
|
||||
};
|
||||
const schema = (full: boolean) =>
|
||||
new Schema([new Field("id", new Int64(), true), field(full)]);
|
||||
const writes: { path: string; schema: Schema }[] = [];
|
||||
let storedWide = wide;
|
||||
const describe = () => ({
|
||||
version: writes.length + 1,
|
||||
schema: {
|
||||
fields: [
|
||||
{ name: "id", type: { type: "int64" }, nullable: true },
|
||||
{
|
||||
name: "image",
|
||||
nullable: true,
|
||||
metadata: { "ARROW:extension:name": "lance.blob.v2" },
|
||||
type: {
|
||||
type: "struct",
|
||||
fields: [
|
||||
{
|
||||
name: "data",
|
||||
nullable: true,
|
||||
type: { type: "large_binary" },
|
||||
},
|
||||
{ name: "uri", nullable: true, type: { type: "string" } },
|
||||
...(storedWide
|
||||
? ["position", "size"].map((name) => ({
|
||||
name,
|
||||
nullable: true,
|
||||
type: { type: "uint64" },
|
||||
}))
|
||||
: []),
|
||||
],
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
});
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
if (req.url?.includes("/describe/")) {
|
||||
res
|
||||
.writeHead(200, { "content-type": "application/json" })
|
||||
.end(JSON.stringify(describe()));
|
||||
return;
|
||||
}
|
||||
const chunks: Buffer[] = [];
|
||||
// Failed streaming input aborts before end; only complete bodies commit.
|
||||
req.on("error", () => {
|
||||
/* Aborted inputs must not commit. */
|
||||
});
|
||||
req.on("data", (chunk) => chunks.push(chunk));
|
||||
req.on("end", () => {
|
||||
const reader = RecordBatchReader.from(Buffer.concat(chunks)).open();
|
||||
// Only read the schema: Rust separately checks compressed wire values.
|
||||
writes.push({ path: req.url ?? "", schema: reader.schema });
|
||||
storedWide =
|
||||
(reader.schema.fields[1].type as Struct).children.length === 4;
|
||||
res.writeHead(200, { "content-type": "application/json" }).end(
|
||||
JSON.stringify({
|
||||
version: writes.length + 1,
|
||||
// biome-ignore lint/style/useNamingConvention: server wire format
|
||||
num_inserted_rows: 1,
|
||||
}),
|
||||
);
|
||||
});
|
||||
},
|
||||
async (db) => {
|
||||
const table = await db.openTable("test");
|
||||
const input = makeArrowTable(
|
||||
[{ id: 1n, image: Buffer.from("payload") }],
|
||||
{ schema: schema(!wide) },
|
||||
);
|
||||
await table.add(input);
|
||||
await table
|
||||
.mergeInsert("id")
|
||||
.whenMatchedUpdateAll()
|
||||
.whenNotMatchedInsertAll()
|
||||
.execute(input);
|
||||
expect(writes).toHaveLength(2);
|
||||
for (const write of writes)
|
||||
expect(write.schema).toEqual(schema(wide));
|
||||
expect(writes[0].path).toContain("/insert/");
|
||||
expect(writes[1].path).toContain("/merge_insert/");
|
||||
if (!wide) {
|
||||
const invalid = makeArrowTable(
|
||||
[
|
||||
{
|
||||
id: 1n,
|
||||
image: {
|
||||
uri: "s3://registered/payload",
|
||||
position: 4,
|
||||
size: 6,
|
||||
},
|
||||
},
|
||||
],
|
||||
{ schema: schema(true) },
|
||||
);
|
||||
await expect(table.add(invalid)).rejects.toThrow(
|
||||
/cannot store byte ranges/,
|
||||
);
|
||||
await expect(
|
||||
table
|
||||
.mergeInsert("id")
|
||||
.whenMatchedUpdateAll()
|
||||
.whenNotMatchedInsertAll()
|
||||
.execute(invalid),
|
||||
).rejects.toThrow(/cannot store byte ranges/);
|
||||
expect(writes).toHaveLength(2);
|
||||
}
|
||||
await table.add(input, { mode: "overwrite" });
|
||||
expect(writes).toHaveLength(3);
|
||||
expect(writes[2].path).toContain("mode=overwrite");
|
||||
expect(writes[2].schema).toEqual(schema(!wide));
|
||||
expect(await table.schema()).toEqual(schema(!wide));
|
||||
},
|
||||
);
|
||||
},
|
||||
);
|
||||
|
||||
it("reports unsupported operations with backend-neutral errors", async () => {
|
||||
await withMockDatabase(
|
||||
(req, res) => {
|
||||
|
||||
@@ -2583,6 +2583,16 @@ describe("when dealing with versioning", () => {
|
||||
});
|
||||
});
|
||||
|
||||
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("when dealing with blob columns", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
beforeEach(() => {
|
||||
@@ -2592,6 +2602,89 @@ describe("when dealing with blob columns", () => {
|
||||
tmpDir.removeCallback();
|
||||
});
|
||||
|
||||
it.each([false, true])(
|
||||
"normalizes Arrow blob layouts for append and merge (range target: %s)",
|
||||
async (wide) => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const schema = (full: boolean) =>
|
||||
new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
full ? blob("image") : legacyBlob("image"),
|
||||
]);
|
||||
const table = await db.createTable(
|
||||
"layouts",
|
||||
[{ id: 1n, image: Buffer.from("old") }],
|
||||
{ schema: schema(wide) },
|
||||
);
|
||||
await table.add(
|
||||
makeArrowTable([{ id: 2n, image: Buffer.from("appended") }], {
|
||||
schema: schema(!wide),
|
||||
}),
|
||||
);
|
||||
await table
|
||||
.mergeInsert("id")
|
||||
.whenMatchedUpdateAll()
|
||||
.whenNotMatchedInsertAll()
|
||||
.execute(
|
||||
makeArrowTable(
|
||||
[
|
||||
{ id: 1n, image: Buffer.from("updated") },
|
||||
{ id: 3n, image: Buffer.from("inserted") },
|
||||
],
|
||||
{ schema: schema(!wide) },
|
||||
),
|
||||
);
|
||||
expect(await table.schema()).toEqual(schema(wide));
|
||||
const rows = await table.query().withRowId().toArray();
|
||||
rows.sort((a, b) => Number(a.id - b.id));
|
||||
expect(
|
||||
await table.fetchBlobs(
|
||||
"image",
|
||||
rows.map((row) => row._rowid as bigint),
|
||||
),
|
||||
).toEqual([
|
||||
Buffer.from("updated"),
|
||||
Buffer.from("appended"),
|
||||
Buffer.from("inserted"),
|
||||
]);
|
||||
if (!wide) {
|
||||
const bad = makeArrowTable(
|
||||
[
|
||||
{
|
||||
id: 1n,
|
||||
image: { uri: "s3://registered/file", position: 0, size: 6 },
|
||||
},
|
||||
],
|
||||
{ schema: schema(true) },
|
||||
);
|
||||
const version = await table.version();
|
||||
await expect(table.add(bad)).rejects.toThrow(
|
||||
/cannot store byte ranges/,
|
||||
);
|
||||
await expect(
|
||||
table
|
||||
.mergeInsert("id")
|
||||
.whenMatchedUpdateAll()
|
||||
.whenNotMatchedInsertAll()
|
||||
.execute(bad),
|
||||
).rejects.toThrow(/cannot store byte ranges/);
|
||||
expect(await table.version()).toBe(version);
|
||||
expect(await table.countRows()).toBe(3);
|
||||
}
|
||||
const replacement = makeArrowTable(
|
||||
[{ id: 4n, image: Buffer.from("replacement") }],
|
||||
{ schema: schema(!wide) },
|
||||
);
|
||||
await table.add(replacement, { mode: "overwrite" });
|
||||
expect(await table.schema()).toEqual(schema(!wide));
|
||||
expect(await table.countRows()).toBe(1);
|
||||
const [row] = await table.query().withRowId().toArray();
|
||||
expect(await table.fetchBlobs("image", [row._rowid as bigint])).toEqual([
|
||||
Buffer.from("replacement"),
|
||||
]);
|
||||
},
|
||||
);
|
||||
|
||||
it("rejects unsupported blob values in createTable and add", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const schema = new Schema([
|
||||
@@ -2668,6 +2761,100 @@ describe("when dealing with blob columns", () => {
|
||||
expect(bytes[0]).toEqual(payload);
|
||||
});
|
||||
|
||||
it.each([0, 4])(
|
||||
"fetches the exact external range at position %s after reopening",
|
||||
async (position) => {
|
||||
const { db, uri, payload } = await openExternalBlobTable();
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int64(), true),
|
||||
blob("payload"),
|
||||
]);
|
||||
const rows: { id: bigint; payload: BlobInput }[] = [
|
||||
{ id: 1n, payload: Buffer.from("inline") },
|
||||
{ id: 2n, payload: { uri: new URL(uri), position, size: 6 } },
|
||||
{ id: 3n, payload: uri },
|
||||
{ id: 4n, payload: null },
|
||||
];
|
||||
const table = await db.createTable("ranges", rows.slice(0, 1), {
|
||||
schema,
|
||||
});
|
||||
await table.add(rows.slice(1), { allowExternalBlobOutsideBases: true });
|
||||
table.close();
|
||||
const reopened = await db.openTable("ranges");
|
||||
const results = await reopened.query().withRowId().toArray();
|
||||
results.sort((a, b) => Number(a.id - b.id));
|
||||
expect(results[1].payload.position).toBe(BigInt(position));
|
||||
expect(results[1].payload.size).toBe(6n);
|
||||
expect(
|
||||
await reopened.fetchBlobs(
|
||||
"payload",
|
||||
results.map((r) => r._rowid as bigint),
|
||||
),
|
||||
).toEqual([
|
||||
Buffer.from("inline"),
|
||||
payload.subarray(position, position + 6),
|
||||
payload,
|
||||
null,
|
||||
]);
|
||||
},
|
||||
);
|
||||
|
||||
it.each([0, 4])(
|
||||
"rejects ranges at position %s on reopened minimal tables without writing",
|
||||
async (position) => {
|
||||
const { db, uri, payload } = await openExternalBlobTable();
|
||||
const table = await db.createEmptyTable(
|
||||
"legacy_blobs",
|
||||
new Schema([new Field("id", new Int64(), true), legacyBlob("payload")]),
|
||||
);
|
||||
table.close();
|
||||
const reopened = await db.openTable("legacy_blobs");
|
||||
await expect(
|
||||
reopened.add(
|
||||
[
|
||||
{ id: 1n, payload: Buffer.from("inline") },
|
||||
{ id: 2n, payload: { uri, position, size: 6 } },
|
||||
],
|
||||
{ allowExternalBlobOutsideBases: true },
|
||||
),
|
||||
).rejects.toThrow(/storage schema cannot store byte ranges/);
|
||||
expect(await reopened.countRows()).toBe(0);
|
||||
await reopened.add(
|
||||
[
|
||||
{ id: 3n, payload: uri },
|
||||
{ id: 4n, payload: Buffer.from("inline") },
|
||||
],
|
||||
{ allowExternalBlobOutsideBases: true },
|
||||
);
|
||||
const rows = await reopened.query().withRowId().toArray();
|
||||
rows.sort((a, b) => Number(a.id - b.id));
|
||||
expect(
|
||||
await reopened.fetchBlobs(
|
||||
"payload",
|
||||
rows.map((r) => r._rowid as bigint),
|
||||
),
|
||||
).toEqual([payload, Buffer.from("inline")]);
|
||||
},
|
||||
);
|
||||
|
||||
it.each([{ position: 0 }, { size: 6 }, { position: 0, size: 0 }])(
|
||||
"rejects partial or empty ranges without writing: %p",
|
||||
async (range) => {
|
||||
const { db, uri } = await openExternalBlobTable();
|
||||
const table = await db.createEmptyTable(
|
||||
"invalid_ranges",
|
||||
new Schema([blob("payload")]),
|
||||
);
|
||||
await expect(
|
||||
table.add(
|
||||
[{ payload: Buffer.from("inline") }, { payload: { uri, ...range } }],
|
||||
{ allowExternalBlobOutsideBases: true },
|
||||
),
|
||||
).rejects.toThrow(/both position and size|size must be positive/);
|
||||
expect(await table.countRows()).toBe(0);
|
||||
},
|
||||
);
|
||||
|
||||
it("applies the external blob opt-in only to the current add", async () => {
|
||||
const { table, uri } = await openExternalBlobTable();
|
||||
|
||||
|
||||
+27
-7
@@ -456,14 +456,29 @@ export function makeArrowTable(
|
||||
|
||||
if (schema !== undefined) {
|
||||
// Validate blob values up front and give every one the full
|
||||
// `{ data, uri }` shape, so inference sees the same struct whether a row
|
||||
// storage shape, so inference sees the same struct whether a row
|
||||
// passed bytes, a URI, or a partial struct.
|
||||
data = mapBlobInputs(data, schema, (value, field, row) => {
|
||||
data = mapBlobInputs(data, schema, (value, field, row, storageField) => {
|
||||
if (value === undefined) {
|
||||
return value;
|
||||
}
|
||||
try {
|
||||
return coerceBlobValue(value);
|
||||
const normalized = coerceBlobValue(value, storageField);
|
||||
if (
|
||||
normalized === null ||
|
||||
!(storageField.type instanceof Struct) ||
|
||||
storageField.type.children.length === 2
|
||||
) {
|
||||
return normalized;
|
||||
}
|
||||
// Keep even unused range children during schema inference. Otherwise
|
||||
// creating a range-capable table with inline data would lose them.
|
||||
return Object.fromEntries(
|
||||
storageField.type.children.map((child) => [
|
||||
child.name,
|
||||
normalized[child.name as keyof typeof normalized] ?? null,
|
||||
]),
|
||||
);
|
||||
} catch (e) {
|
||||
throw new Error(
|
||||
`Invalid value for blob field ${field} at row ${row}: ${(e as Error).message}`,
|
||||
@@ -533,7 +548,12 @@ function containsBlobField(field: Field): boolean {
|
||||
);
|
||||
}
|
||||
|
||||
type BlobInputVisitor = (value: unknown, field: string, row: number) => unknown;
|
||||
type BlobInputVisitor = (
|
||||
value: unknown,
|
||||
field: string,
|
||||
row: number,
|
||||
storageField: Field,
|
||||
) => unknown;
|
||||
|
||||
/**
|
||||
* Calls `visit` on every value that lands in a blob column of `schema`,
|
||||
@@ -592,7 +612,7 @@ function mapBlobField(
|
||||
visit: BlobInputVisitor,
|
||||
): unknown {
|
||||
if (isBlobField(field)) {
|
||||
return visit(value, label, row);
|
||||
return visit(value, label, row, field);
|
||||
}
|
||||
if (field.type instanceof Struct && isObject(value)) {
|
||||
return mapBlobFields(value, field.type.children, `${label}.`, row, visit);
|
||||
@@ -697,11 +717,11 @@ function transposeData(
|
||||
const valuesPath = [...path, field.name];
|
||||
if (isBlobField(field) && field.type instanceof Struct) {
|
||||
const blobRows = data.map((datum) =>
|
||||
coerceBlobValue(valueAtPath(datum, valuesPath)),
|
||||
coerceBlobValue(valueAtPath(datum, valuesPath), field),
|
||||
);
|
||||
const childVectors = field.type.children.map((child) => {
|
||||
const values = blobRows.map((row) =>
|
||||
row == null ? null : (row[child.name as "data" | "uri"] ?? null),
|
||||
row == null ? null : (row[child.name as keyof typeof row] ?? null),
|
||||
);
|
||||
return makeVector(values, child.type, undefined, child.nullable);
|
||||
});
|
||||
|
||||
+91
-8
@@ -1,7 +1,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { Field, LargeBinary, Struct, Utf8 } from "apache-arrow";
|
||||
import { Field, LargeBinary, Struct, Uint64, Utf8 } from "apache-arrow";
|
||||
import { BlobFile as NativeBlobFile } from "./native";
|
||||
|
||||
const BLOB_V2_EXTENSION_NAME = "lance.blob.v2";
|
||||
@@ -28,6 +28,9 @@ export type BlobUri = string | URL;
|
||||
*
|
||||
* Either inline bytes, a URI pointing at external bytes, or a struct that sets
|
||||
* exactly one of `data` or `uri`. `null` and `undefined` write a null blob.
|
||||
* External byte ranges require both
|
||||
* `position` (non-negative) and `size` (positive). Use safe integer numbers or
|
||||
* `bigint` values for these unsigned 64-bit byte counts.
|
||||
*
|
||||
* @example
|
||||
* ```ts
|
||||
@@ -45,8 +48,14 @@ export type BlobUri = string | URL;
|
||||
export type BlobInput =
|
||||
| BlobData
|
||||
| BlobUri
|
||||
| { data: BlobData; uri?: null }
|
||||
| { data?: null; uri: BlobUri }
|
||||
| { data: BlobData; uri?: null; position?: null; size?: null }
|
||||
| { data?: null; uri: BlobUri; position?: null; size?: null }
|
||||
| {
|
||||
data?: null;
|
||||
uri: BlobUri;
|
||||
position: number | bigint;
|
||||
size: number | bigint;
|
||||
}
|
||||
| null
|
||||
| undefined;
|
||||
|
||||
@@ -54,6 +63,8 @@ export type BlobInput =
|
||||
export type BlobValue = {
|
||||
data: Buffer | Uint8Array | null;
|
||||
uri: string | null;
|
||||
position?: bigint;
|
||||
size?: bigint;
|
||||
};
|
||||
|
||||
export type BlobOptions = {
|
||||
@@ -78,6 +89,15 @@ export type BlobOptions = {
|
||||
|
||||
/**
|
||||
* Declares a `lance.blob.v2` column.
|
||||
* Uses nullable `data`, `uri`, `position`, and `size` fields for inline bytes,
|
||||
* whole external files, and external byte ranges. Set both `position` and
|
||||
* positive `size` for a range, or omit both for a whole file.
|
||||
*
|
||||
* Existing two-field columns remain writable. Appends and merge inserts fill
|
||||
* missing range fields with nulls for four-field targets and drop all-null
|
||||
* range fields for two-field targets. Non-null ranges require a four-field
|
||||
* target. Overwrite with Arrow input replaces the table's schema; row objects
|
||||
* are converted using the current table schema.
|
||||
*
|
||||
* Query results are descriptors, not payload bytes. Use {@link Table.fetchBlobs}
|
||||
* or {@link Table.fetchBlobFiles} to read bytes.
|
||||
@@ -142,6 +162,8 @@ export function blob(name: string, options: BlobOptions = {}): Field {
|
||||
new Struct([
|
||||
new Field("data", new LargeBinary(), true),
|
||||
new Field("uri", new Utf8(), true),
|
||||
new Field("position", new Uint64(), true),
|
||||
new Field("size", new Uint64(), true),
|
||||
]),
|
||||
options.nullable ?? true,
|
||||
metadata,
|
||||
@@ -317,6 +339,8 @@ export function normalizeBlobInput(value: unknown): unknown {
|
||||
...value,
|
||||
data: value.data,
|
||||
uri: value.uri,
|
||||
position: value.position,
|
||||
size: value.size,
|
||||
};
|
||||
let changed = false;
|
||||
if (value.data instanceof ArrayBuffer) {
|
||||
@@ -355,10 +379,19 @@ export function withBlobBytes(value: unknown, bytes: Uint8Array): unknown {
|
||||
return bytes;
|
||||
}
|
||||
const record = value as Record<string, unknown>;
|
||||
return { ...record, uri: record.uri, data: bytes };
|
||||
return {
|
||||
...record,
|
||||
uri: record.uri,
|
||||
position: record.position,
|
||||
size: record.size,
|
||||
data: bytes,
|
||||
};
|
||||
}
|
||||
|
||||
export function coerceBlobValue(input: unknown): BlobValue | null {
|
||||
export function coerceBlobValue(
|
||||
input: unknown,
|
||||
field?: Field,
|
||||
): BlobValue | null {
|
||||
const value = normalizeBlobInput(input);
|
||||
if (value == null) {
|
||||
return null;
|
||||
@@ -385,11 +418,11 @@ export function coerceBlobValue(input: unknown): BlobValue | null {
|
||||
);
|
||||
}
|
||||
const unknownKeys = Object.keys(record).filter(
|
||||
(key) => key !== "data" && key !== "uri",
|
||||
(key) => !["data", "uri", "position", "size"].includes(key),
|
||||
);
|
||||
if (unknownKeys.length > 0) {
|
||||
throw new Error(
|
||||
`Blob struct values only support 'data' and 'uri', got ${unknownKeys
|
||||
`Blob struct values only support 'data', 'uri', 'position', and 'size', got ${unknownKeys
|
||||
.map((key) => `'${key}'`)
|
||||
.join(", ")}`,
|
||||
);
|
||||
@@ -416,13 +449,63 @@ export function coerceBlobValue(input: unknown): BlobValue | null {
|
||||
"Blob struct values must set exactly one of 'data' or 'uri'",
|
||||
);
|
||||
}
|
||||
const hasPosition = record.position != null;
|
||||
const hasSize = record.size != null;
|
||||
if (hasPosition !== hasSize) {
|
||||
throw new Error(
|
||||
"External blob must set both position and size, or neither",
|
||||
);
|
||||
}
|
||||
if (hasPosition) {
|
||||
if (uriValue == null) {
|
||||
throw new Error("External blob byte ranges require a uri");
|
||||
}
|
||||
const position = blobRangeInteger(record.position, "position", 0n);
|
||||
const size = blobRangeInteger(record.size, "size", 1n);
|
||||
if (position + size > 0xffffffffffffffffn) {
|
||||
throw new Error("External blob range overflows uint64");
|
||||
}
|
||||
if (
|
||||
field !== undefined &&
|
||||
(!(field.type instanceof Struct) ||
|
||||
["position", "size"].some(
|
||||
(name) =>
|
||||
!field.type.children.some((child: Field) => child.name === name),
|
||||
))
|
||||
) {
|
||||
throw new Error(
|
||||
"Blob storage schema cannot store byte ranges. Use a four-field blob() schema to store position and size",
|
||||
);
|
||||
}
|
||||
return { data: null, uri: uriValue, position, size };
|
||||
}
|
||||
return { data: bytes, uri: uriValue };
|
||||
}
|
||||
throw new Error(
|
||||
"Blob column values must be Buffer, Uint8Array, ArrayBuffer, Blob, a URI string or URL, null, or { data?, uri? }",
|
||||
"Blob column values must be Buffer, Uint8Array, ArrayBuffer, Blob, a URI string or URL, null, or { data?, uri?, position?, size? }",
|
||||
);
|
||||
}
|
||||
|
||||
function blobRangeInteger(
|
||||
value: unknown,
|
||||
name: string,
|
||||
minimum: bigint,
|
||||
): bigint {
|
||||
if (
|
||||
typeof value !== "bigint" &&
|
||||
(typeof value !== "number" || !Number.isSafeInteger(value))
|
||||
) {
|
||||
throw new Error(`Blob ${name} must be a safe integer or bigint`);
|
||||
}
|
||||
const integer = BigInt(value);
|
||||
if (integer < minimum || integer > 0xffffffffffffffffn) {
|
||||
throw new Error(
|
||||
`Blob ${name} must be ${minimum === 0n ? "non-negative" : "positive"} and fit in uint64`,
|
||||
);
|
||||
}
|
||||
return integer;
|
||||
}
|
||||
|
||||
function isBlobLike(value: unknown): value is Blob {
|
||||
return typeof Blob !== "undefined" && value instanceof Blob;
|
||||
}
|
||||
|
||||
@@ -588,7 +588,7 @@ def _cast_to_target_schema(
|
||||
def _preserve_blob_v2_fields(
|
||||
inferred_schema: pa.Schema, existing_schema: pa.Schema
|
||||
) -> pa.Schema:
|
||||
"""Keep existing blob fields when an overwrite infers a replacement schema."""
|
||||
"""Keep inferred blob semantics without overriding explicit replacement types."""
|
||||
existing_fields = {field.name: field for field in existing_schema}
|
||||
fields = [
|
||||
_preserve_blob_v2_field(field, existing_fields[field.name])
|
||||
@@ -600,6 +600,9 @@ def _preserve_blob_v2_fields(
|
||||
|
||||
|
||||
def _preserve_blob_v2_field(inferred: pa.Field, existing: pa.Field) -> pa.Field:
|
||||
if is_blob_v2_field(inferred):
|
||||
# An explicitly declared blob layout is part of the replacement schema.
|
||||
return inferred
|
||||
if is_blob_v2_field(existing):
|
||||
return existing
|
||||
|
||||
@@ -637,12 +640,10 @@ def _preserve_blob_v2_field(inferred: pa.Field, existing: pa.Field) -> pa.Field:
|
||||
def _coerce_blob_write_columns(
|
||||
batch: pa.RecordBatch, target_schema: pa.Schema
|
||||
) -> pa.RecordBatch:
|
||||
"""Materialize blob storage structs before the stream leaves Python.
|
||||
"""Materialize Python bytes and URIs as blob structs before Arrow casting.
|
||||
|
||||
merge_insert requires its source reader to already match the table's
|
||||
physical schema. Unlike add and insert, it does not pass through
|
||||
LanceDB's Rust blob coercion, so preserving binary input here would
|
||||
reach Lance as binary and fail the schema check.
|
||||
Existing blob structs retain their input layout. Rust normalizes them to
|
||||
the stored schema and rejects populated ranges on two-field targets.
|
||||
"""
|
||||
columns = []
|
||||
fields = []
|
||||
@@ -973,6 +974,17 @@ def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
|
||||
"(Utf8, LargeUtf8, Utf8View), or a Struct with a "
|
||||
"'data' or 'uri' child"
|
||||
)
|
||||
if pa.types.is_struct(source_type) or (
|
||||
isinstance(source_type, pa.ExtensionType)
|
||||
and (
|
||||
not isinstance(target_field.type, pa.ExtensionType)
|
||||
or type(source_type) is type(target_field.type)
|
||||
)
|
||||
):
|
||||
# Keep the source layout until Rust can fill missing range fields or
|
||||
# validate every batch before dropping all-null range fields. Arrow
|
||||
# casting here would reject differing layouts before that validation.
|
||||
return field.with_metadata(target_field.metadata)
|
||||
# LanceDB exposes stored JSON columns as lance.json (JSONB-backed LargeBinary), but
|
||||
# casting the input to that storage type here merely relabels the raw JSON bytes as
|
||||
# JSONB. Lance must see arrow.json so it can perform the JSONB encoding.
|
||||
@@ -1832,8 +1844,8 @@ class Table(ABC):
|
||||
- pyarrow.Table or pyarrow.RecordBatch
|
||||
mode: str
|
||||
The mode to use when writing the data. Valid values are
|
||||
"append" and "overwrite". Overwrite infers a replacement schema
|
||||
but preserves existing blob v2 fields present in the new data.
|
||||
"append" and "overwrite". Overwrite adopts explicitly declared blob
|
||||
layouts, while inferred bytes retain the existing blob v2 fields.
|
||||
on_bad_vectors: str, default "error"
|
||||
What to do if any of the vectors are not the same size or contains NaNs.
|
||||
One of "error", "drop", "fill".
|
||||
@@ -4111,8 +4123,8 @@ class LanceTable(Table):
|
||||
The data to insert into the table.
|
||||
mode: str
|
||||
The mode to use when writing the data. Valid values are
|
||||
"append" and "overwrite". Overwrite infers a replacement schema
|
||||
but preserves existing blob v2 fields present in the new data.
|
||||
"append" and "overwrite". Overwrite adopts explicitly declared blob
|
||||
layouts, while inferred bytes retain the existing blob v2 fields.
|
||||
on_bad_vectors: str, default "error"
|
||||
What to do if any of the vectors are not the same size or contains NaNs.
|
||||
One of "error", "drop", "fill", "null".
|
||||
@@ -6145,8 +6157,8 @@ class AsyncTable:
|
||||
- pyarrow.Table or pyarrow.RecordBatch
|
||||
mode: str
|
||||
The mode to use when writing the data. Valid values are
|
||||
"append" and "overwrite". Overwrite infers a replacement schema
|
||||
but preserves existing blob v2 fields present in the new data.
|
||||
"append" and "overwrite". Overwrite adopts explicitly declared blob
|
||||
layouts, while inferred bytes retain the existing blob v2 fields.
|
||||
on_bad_vectors: str, default "error"
|
||||
What to do if any of the vectors are not the same size or contains NaNs.
|
||||
One of "error", "drop", "fill", "null".
|
||||
|
||||
@@ -63,6 +63,25 @@ def _blob_array(name, values):
|
||||
return pa.ExtensionArray.from_storage(blob_type, storage)
|
||||
|
||||
|
||||
def _blob_layout_data(rows, with_ranges, *, as_extension=True):
|
||||
field = lancedb.blob("image")
|
||||
storage_type = field.type.storage_type
|
||||
if not with_ranges or not as_extension:
|
||||
storage_type = pa.struct(list(storage_type)[: (4 if with_ranges else 2)])
|
||||
field = pa.field(
|
||||
"image",
|
||||
storage_type,
|
||||
metadata={"ARROW:extension:name": "lance.blob.v2"},
|
||||
)
|
||||
images = pa.array([row["image"] for row in rows], type=storage_type)
|
||||
if isinstance(field.type, pa.ExtensionType):
|
||||
images = pa.ExtensionArray.from_storage(field.type, images)
|
||||
return pa.Table.from_arrays(
|
||||
[pa.array([row["id"] for row in rows], type=pa.int64()), images],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), field]),
|
||||
)
|
||||
|
||||
|
||||
def _row_ids_by_id(table):
|
||||
hits = table.search().with_row_id(True).limit(1000).to_arrow()
|
||||
assert "_rowid" in hits.column_names
|
||||
@@ -998,6 +1017,133 @@ def test_merge_insert_writes_python_bytes():
|
||||
assert blobs.to_pylist() == [b"updated", b"inserted"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
@pytest.mark.parametrize("input_type", ["extension", "struct", "reader"])
|
||||
@pytest.mark.parametrize("target_with_ranges", [False, True])
|
||||
async def test_merge_insert_normalizes_blob_layouts(
|
||||
use_async, input_type, target_with_ranges
|
||||
):
|
||||
data = _blob_layout_data(
|
||||
[{"id": 1, "image": {"data": b"before"}}], target_with_ranges
|
||||
)
|
||||
table = lancedb.connect("memory:///").create_table("merge_layouts", data=data)
|
||||
original_schema = table.schema
|
||||
updates = _blob_layout_data(
|
||||
[
|
||||
{"id": 1, "image": {"data": b"updated"}},
|
||||
{"id": 2, "image": {"data": b"inserted"}},
|
||||
{"id": 3, "image": None},
|
||||
],
|
||||
not target_with_ranges,
|
||||
as_extension=input_type != "struct",
|
||||
)
|
||||
if input_type == "reader":
|
||||
updates = updates.to_reader(max_chunksize=1)
|
||||
writer = table._table if use_async else table
|
||||
result = (
|
||||
writer.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(updates)
|
||||
)
|
||||
if use_async:
|
||||
result = await result
|
||||
|
||||
assert result.num_updated_rows == 1
|
||||
assert result.num_inserted_rows == 2
|
||||
assert table.schema == original_schema
|
||||
by_id = _row_ids_by_id(table)
|
||||
assert table.fetch_blobs("image", [by_id[i] for i in (1, 2, 3)]).to_pylist() == [
|
||||
b"updated",
|
||||
b"inserted",
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"byte_range", [{"position": 0}, {"size": 6}, {"position": 4, "size": 6}]
|
||||
)
|
||||
async def test_merge_insert_legacy_blob_rejects_ranges_atomically(
|
||||
use_async, byte_range
|
||||
):
|
||||
data = _blob_layout_data([{"id": 1, "image": {"data": b"before"}}], False)
|
||||
table = lancedb.connect("memory:///").create_table("reject_ranges", data=data)
|
||||
original_schema, original_version = table.schema, table.version
|
||||
updates = _blob_layout_data(
|
||||
[
|
||||
{"id": 1, "image": {"data": b"updated"}},
|
||||
{"id": 2, "image": {"uri": "file:///payload", **byte_range}},
|
||||
],
|
||||
True,
|
||||
).to_reader(max_chunksize=1)
|
||||
writer = table._table if use_async else table
|
||||
with pytest.raises(RuntimeError, match="cannot store byte ranges"):
|
||||
result = (
|
||||
writer.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(updates)
|
||||
)
|
||||
if use_async:
|
||||
await result
|
||||
|
||||
assert table.version == original_version
|
||||
assert table.schema == original_schema
|
||||
assert table.count_rows() == 1
|
||||
assert table.fetch_blobs("image", [_row_ids_by_id(table)[1]]).to_pylist() == [
|
||||
b"before"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
@pytest.mark.parametrize("input_type", ["extension", "struct", "reader"])
|
||||
@pytest.mark.parametrize("source_with_ranges", [False, True])
|
||||
async def test_add_overwrite_adopts_explicit_blob_layout(
|
||||
tmp_path, use_async, input_type, source_with_ranges
|
||||
):
|
||||
data = _blob_layout_data(
|
||||
[{"id": 1, "image": {"data": b"before"}}], not source_with_ranges
|
||||
)
|
||||
table = lancedb.connect("memory:///").create_table("overwrite_layouts", data=data)
|
||||
payload = b"0123456789abcdef"
|
||||
path = tmp_path / "payload.bin"
|
||||
path.write_bytes(payload)
|
||||
value = (
|
||||
{"uri": path.as_uri(), "position": 4, "size": 6}
|
||||
if source_with_ranges
|
||||
else {"data": b"updated"}
|
||||
)
|
||||
updates = _blob_layout_data(
|
||||
[{"id": 2, "image": value}, {"id": 3, "image": None}],
|
||||
source_with_ranges,
|
||||
as_extension=input_type != "struct",
|
||||
)
|
||||
if input_type == "reader":
|
||||
updates = updates.to_reader(max_chunksize=1)
|
||||
writer = table._table if use_async else table
|
||||
result = writer.add(
|
||||
updates, mode="overwrite", allow_external_blob_outside_bases=True
|
||||
)
|
||||
if use_async:
|
||||
await result
|
||||
|
||||
blob_type = table.schema.field("image").type
|
||||
if isinstance(blob_type, pa.ExtensionType):
|
||||
blob_type = blob_type.storage_type
|
||||
assert blob_type.num_fields == (4 if source_with_ranges else 2)
|
||||
assert table.blob_columns() == ["image"]
|
||||
assert table.count_rows() == 2
|
||||
by_id = _row_ids_by_id(table)
|
||||
assert table.fetch_blobs("image", [by_id[2], by_id[3]]).to_pylist() == [
|
||||
payload[4:10] if source_with_ranges else b"updated",
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("write", ["merge_insert", "create_table"])
|
||||
def test_uri_string_blob_write_reaches_external_base_validation(tmp_path, write):
|
||||
blob_path = tmp_path / "payload.bin"
|
||||
|
||||
@@ -3254,6 +3254,89 @@ BLOB_DESCRIBE_RESPONSE = {
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_async", [False, True])
|
||||
@pytest.mark.parametrize("mode", ["merge", "overwrite"])
|
||||
@pytest.mark.parametrize("target_with_ranges", [False, True])
|
||||
async def test_remote_writes_normalize_blob_layouts(
|
||||
use_async, mode, target_with_ranges
|
||||
):
|
||||
description = json.loads(json.dumps(BLOB_DESCRIBE_RESPONSE))
|
||||
if target_with_ranges:
|
||||
description["schema"]["fields"][1]["type"]["fields"].extend(
|
||||
{"name": name, "type": {"type": "uint64"}, "nullable": True}
|
||||
for name in ("position", "size")
|
||||
)
|
||||
received = []
|
||||
|
||||
def handler(request):
|
||||
if request.path == "/v1/table/test/describe/":
|
||||
response = description
|
||||
else:
|
||||
expected = "merge_insert" if mode == "merge" else "insert"
|
||||
assert urlparse(request.path).path == f"/v1/table/test/{expected}/"
|
||||
if request.headers.get("Transfer-Encoding") == "chunked":
|
||||
body = bytearray()
|
||||
while size := int(request.rfile.readline().split(b";")[0], 16):
|
||||
body.extend(request.rfile.read(size))
|
||||
assert request.rfile.read(2) == b"\r\n"
|
||||
assert request.rfile.read(2) == b"\r\n"
|
||||
else:
|
||||
body = request.rfile.read(int(request.headers["Content-Length"]))
|
||||
received.append(pa.ipc.open_stream(body).read_all())
|
||||
response = {
|
||||
"version": 2,
|
||||
"num_updated_rows": 0,
|
||||
"num_inserted_rows": 2,
|
||||
"num_deleted_rows": 0,
|
||||
}
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.send_header("phalanx-version", "0.5.0")
|
||||
request.end_headers()
|
||||
request.wfile.write(json.dumps(response).encode())
|
||||
|
||||
field = lancedb.blob("image")
|
||||
storage_type = field.type.storage_type
|
||||
if target_with_ranges:
|
||||
storage_type = pa.struct(list(storage_type)[:2])
|
||||
field = pa.field(
|
||||
"image",
|
||||
storage_type,
|
||||
metadata={"ARROW:extension:name": "lance.blob.v2"},
|
||||
)
|
||||
images = pa.array([{"data": b"payload"}, None], type=storage_type)
|
||||
if isinstance(field.type, pa.ExtensionType):
|
||||
images = pa.ExtensionArray.from_storage(field.type, images)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array([1, 2], type=pa.int64()), images],
|
||||
schema=pa.schema([pa.field("id", pa.int64()), field]),
|
||||
)
|
||||
|
||||
with mock_lancedb_connection(handler) as db:
|
||||
remote_table = db.open_table("test")
|
||||
table = remote_table._table if use_async else remote_table
|
||||
if mode == "merge":
|
||||
result = (
|
||||
table.merge_insert("id").when_not_matched_insert_all().execute(data)
|
||||
)
|
||||
else:
|
||||
result = table.add(data, mode="overwrite")
|
||||
if use_async:
|
||||
await result
|
||||
|
||||
assert len(received) == 1
|
||||
images = received[0]["image"].combine_chunks()
|
||||
if isinstance(images.type, pa.ExtensionType):
|
||||
images = images.storage
|
||||
with_ranges = target_with_ranges if mode == "merge" else not target_with_ranges
|
||||
assert images.type.num_fields == (4 if with_ranges else 2)
|
||||
expected = {"data": b"payload", "uri": None}
|
||||
if with_ranges:
|
||||
expected.update(position=None, size=None)
|
||||
assert images.to_pylist() == [expected, None]
|
||||
|
||||
|
||||
def blob_query_response_table():
|
||||
image_field = pa.field(
|
||||
"image",
|
||||
|
||||
@@ -244,8 +244,34 @@ impl BlobRangeRequest {
|
||||
|
||||
/// Creates an Arrow field for a Lance blob v2 column.
|
||||
///
|
||||
/// `Struct<data, uri>` with the `lance.blob.v2` marker. Same layout Lance
|
||||
/// expects on write.
|
||||
/// `Struct<data, uri, position, size>` with nullable children and the
|
||||
/// `lance.blob.v2` marker. Set both range fields to read a byte range from a
|
||||
/// URI, or leave both null for a whole file. Inline bytes are also supported.
|
||||
///
|
||||
/// # Migration
|
||||
///
|
||||
/// This helper previously returned only `data` and `uri`. Rust callers building
|
||||
/// a `StructArray` from this field must now supply four children: append typed
|
||||
/// null UInt64 arrays for `position` and `size` when no range is needed:
|
||||
///
|
||||
/// ```
|
||||
/// use std::sync::Arc;
|
||||
/// use arrow_array::{ArrayRef, LargeBinaryArray, StringArray, StructArray, UInt64Array};
|
||||
/// use arrow_schema::DataType;
|
||||
/// let field = lancedb::blob("image", true);
|
||||
/// let DataType::Struct(children) = field.data_type() else { unreachable!() };
|
||||
/// let image = StructArray::new(children.clone(), vec![
|
||||
/// Arc::new(LargeBinaryArray::from(vec![Some(b"image".as_slice())])) as ArrayRef,
|
||||
/// Arc::new(StringArray::from(vec![None::<&str>])),
|
||||
/// Arc::new(UInt64Array::new_null(1)),
|
||||
/// Arc::new(UInt64Array::new_null(1)),
|
||||
/// ], None);
|
||||
/// ```
|
||||
///
|
||||
/// Existing two-field tables remain writable: append and merge-insert fill
|
||||
/// missing range fields or drop all-null range fields to match the table.
|
||||
/// Non-null ranges cannot be written to a two-field table. Overwrite replaces
|
||||
/// the table schema with the input schema rather than normalizing its layout.
|
||||
///
|
||||
/// A blob column may be top-level or nested inside a struct or list. Blobs
|
||||
/// nested inside structs use dotted paths (e.g. `info.blob`) in the read APIs.
|
||||
@@ -261,9 +287,21 @@ impl BlobRangeRequest {
|
||||
/// ]);
|
||||
/// ```
|
||||
pub fn blob(name: impl AsRef<str>, nullable: bool) -> Field {
|
||||
lance::blob::blob_field(name.as_ref(), nullable)
|
||||
blob_with_options(name, nullable, BlobOptions::default())
|
||||
}
|
||||
|
||||
static BLOB_STORAGE_TYPE: std::sync::LazyLock<DataType> = std::sync::LazyLock::new(|| {
|
||||
DataType::Struct(
|
||||
vec![
|
||||
Field::new("data", DataType::LargeBinary, true),
|
||||
Field::new("uri", DataType::Utf8, true),
|
||||
Field::new("position", DataType::UInt64, true),
|
||||
Field::new("size", DataType::UInt64, true),
|
||||
]
|
||||
.into(),
|
||||
)
|
||||
});
|
||||
|
||||
/// Options for creating a blob v2 column field with tier thresholds.
|
||||
///
|
||||
/// Re-exported from lance so the thresholds are built exactly the way
|
||||
@@ -293,6 +331,7 @@ pub use lance::blob::BlobFieldOptions as BlobOptions;
|
||||
/// ```
|
||||
pub fn blob_with_options(name: impl AsRef<str>, nullable: bool, options: BlobOptions) -> Field {
|
||||
lance::blob::blob_field_with_options(name.as_ref(), nullable, options)
|
||||
.with_data_type(BLOB_STORAGE_TYPE.clone())
|
||||
}
|
||||
|
||||
/// Returns true if `field` is a blob v2 column.
|
||||
@@ -306,7 +345,7 @@ pub fn is_blob(field: &Field) -> bool {
|
||||
}
|
||||
|
||||
/// Returns true if `field`, or any field nested under it, is a blob v2 column.
|
||||
fn field_tree_has_blob_v2(field: &Field) -> bool {
|
||||
pub(crate) fn field_tree_has_blob_v2(field: &Field) -> bool {
|
||||
if field.is_blob_v2() {
|
||||
return true;
|
||||
}
|
||||
@@ -673,7 +712,21 @@ mod tests {
|
||||
field.metadata().get(ARROW_EXT_NAME_KEY).map(String::as_str),
|
||||
Some("lance.blob.v2")
|
||||
);
|
||||
assert!(matches!(field.data_type(), DataType::Struct(_)));
|
||||
let DataType::Struct(children) = field.data_type() else {
|
||||
panic!("expected struct")
|
||||
};
|
||||
assert_eq!(
|
||||
children
|
||||
.iter()
|
||||
.map(|f| (f.name().as_str(), f.data_type(), f.is_nullable()))
|
||||
.collect::<Vec<_>>(),
|
||||
vec![
|
||||
("data", &DataType::LargeBinary, true),
|
||||
("uri", &DataType::Utf8, true),
|
||||
("position", &DataType::UInt64, true),
|
||||
("size", &DataType::UInt64, true),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -3154,6 +3154,17 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
let rescannable = source.rescannable();
|
||||
let input: Arc<dyn ExecutionPlan> =
|
||||
Arc::new(crate::table::datafusion::scannable_exec::ScannableExec::new(source, None));
|
||||
let table_schema = self.schema().await?;
|
||||
let input = if crate::blob::has_blob_columns(&table_schema) {
|
||||
crate::table::merge::validate_merge_source_fields(
|
||||
input.schema().fields(),
|
||||
table_schema.fields(),
|
||||
"",
|
||||
)?;
|
||||
crate::table::datafusion::cast::cast_to_table_schema(input, &table_schema)?
|
||||
} else {
|
||||
input
|
||||
};
|
||||
let freshness_request = self.snapshot_freshness_headers();
|
||||
|
||||
let mut merge: Arc<dyn ExecutionPlan> = Arc::new(
|
||||
@@ -4499,6 +4510,253 @@ mod tests {
|
||||
assert!(image.column_by_name("data").unwrap().is_null(0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn add_external_blob_ranges_preserves_remote_wire_values() {
|
||||
let field = crate::blob("image", true);
|
||||
let DataType::Struct(children) = field.data_type().clone() else {
|
||||
unreachable!();
|
||||
};
|
||||
let schema = Arc::new(Schema::new(vec![field]));
|
||||
// Registered external bases are enforced by the server, so remote
|
||||
// clients send the URI and range without the local-only opt-in.
|
||||
let uri = "s3://registered/payload.bin";
|
||||
let image = StructArray::new(
|
||||
children,
|
||||
vec![
|
||||
arrow_array::new_null_array(&DataType::LargeBinary, 2),
|
||||
Arc::new(StringArray::from(vec![uri, uri])),
|
||||
Arc::new(arrow_array::UInt64Array::from(vec![0, 4])),
|
||||
Arc::new(arrow_array::UInt64Array::from(vec![6, 6])),
|
||||
],
|
||||
None,
|
||||
);
|
||||
let input = RecordBatch::try_new(schema.clone(), vec![Arc::new(image)]).unwrap();
|
||||
let describe_body = describe_response(&schema).into_bytes();
|
||||
let payload = b"0123456789abcdef";
|
||||
let fetched = RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![Field::new(
|
||||
"image",
|
||||
DataType::LargeBinary,
|
||||
true,
|
||||
)])),
|
||||
vec![Arc::new(LargeBinaryArray::from(vec![
|
||||
&payload[..6],
|
||||
&payload[4..10],
|
||||
]))],
|
||||
)
|
||||
.unwrap();
|
||||
let fetched_body = write_ipc_stream_uncompressed(&fetched);
|
||||
let (sender, receiver) = std::sync::mpsc::channel();
|
||||
let table = Table::new_with_handler_version(
|
||||
"my_table",
|
||||
semver::Version::new(0, 5, 0),
|
||||
move |mut request| match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_body.clone())
|
||||
.unwrap(),
|
||||
"/v1/table/my_table/insert/" => {
|
||||
sender.send(request.body_mut().take().unwrap()).unwrap();
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(br#"{"version": 2}"#.to_vec())
|
||||
.unwrap()
|
||||
}
|
||||
"/v1/table/my_table/fetch_blobs/" => {
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(body["column"], "image");
|
||||
assert_eq!(body["row_ids"], json!([0, 1]));
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE)
|
||||
.body(fetched_body.clone())
|
||||
.unwrap()
|
||||
}
|
||||
path => panic!("unexpected path: {path}"),
|
||||
},
|
||||
);
|
||||
table.add(input).execute().await.unwrap();
|
||||
let body = collect_body(receiver.recv().unwrap()).await;
|
||||
let batch = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(body), None)
|
||||
.unwrap()
|
||||
.next()
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(batch.schema(), schema);
|
||||
let image = batch.column(0).as_struct();
|
||||
let positions = image
|
||||
.column(2)
|
||||
.as_primitive::<arrow::datatypes::UInt64Type>();
|
||||
let sizes = image
|
||||
.column(3)
|
||||
.as_primitive::<arrow::datatypes::UInt64Type>();
|
||||
assert_eq!(positions.values().as_ref(), &[0, 4]);
|
||||
assert_eq!(sizes.values().as_ref(), &[6, 6]);
|
||||
assert_eq!(image.column(1).as_string::<i32>().value(0), uri);
|
||||
assert_eq!(image.column(1).as_string::<i32>().value(1), uri);
|
||||
let bytes = table.fetch_blobs("image", &[0, 1]).await.unwrap();
|
||||
assert_eq!(bytes.value(0), &payload[..6]);
|
||||
assert_eq!(bytes.value(1), &payload[4..10]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_blob_layout_normalization_and_overwrite() {
|
||||
use arrow_array::{ArrayRef, UInt64Array};
|
||||
for wide_target in [false, true] {
|
||||
let field = |wide| {
|
||||
if wide {
|
||||
crate::blob("image", true)
|
||||
} else {
|
||||
lance::blob::blob_field("image", true)
|
||||
}
|
||||
};
|
||||
let target = Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
field(wide_target),
|
||||
]);
|
||||
for operation in ["add", "merge", "overwrite"] {
|
||||
for range in [
|
||||
None,
|
||||
Some((Some(4), None)),
|
||||
Some((None, Some(6))),
|
||||
Some((Some(4), Some(6))),
|
||||
] {
|
||||
if range.is_some() && (wide_target || operation == "overwrite") {
|
||||
continue;
|
||||
}
|
||||
// Range input always uses the full layout. Otherwise use the
|
||||
// opposite layout to exercise both normalization directions.
|
||||
let wide_input = range.is_some() || !wide_target;
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
field(wide_input),
|
||||
]));
|
||||
let DataType::Struct(children) = schema.field(1).data_type().clone() else {
|
||||
unreachable!()
|
||||
};
|
||||
let mut arrays: Vec<ArrayRef> = vec![
|
||||
Arc::new(LargeBinaryArray::from(vec![Some(b"inline".as_slice())])),
|
||||
Arc::new(StringArray::from(vec![None::<&str>])),
|
||||
];
|
||||
if wide_input {
|
||||
arrays.push(Arc::new(UInt64Array::from(vec![range.and_then(|r| r.0)])));
|
||||
arrays.push(Arc::new(UInt64Array::from(vec![range.and_then(|r| r.1)])));
|
||||
}
|
||||
let image = Arc::new(StructArray::new(children, arrays, None));
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![Arc::new(Int64Array::from(vec![1])), image],
|
||||
)
|
||||
.unwrap();
|
||||
let describe = describe_response(&target);
|
||||
let (sender, receiver) = std::sync::mpsc::channel();
|
||||
let table = Table::new_with_handler("my_table", move |mut request| {
|
||||
if request.url().path().ends_with("/describe/") {
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe.clone())
|
||||
.unwrap();
|
||||
}
|
||||
assert!(
|
||||
request.url().path().ends_with("/insert/")
|
||||
|| request.url().path().ends_with("/merge_insert/")
|
||||
);
|
||||
if operation == "overwrite" {
|
||||
assert!(
|
||||
request
|
||||
.url()
|
||||
.query_pairs()
|
||||
.any(|(k, v)| k == "mode" && v == "overwrite")
|
||||
);
|
||||
}
|
||||
sender.send(request.body_mut().take().unwrap()).unwrap();
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version": 2, "num_inserted_rows": 1}"#.to_string())
|
||||
.unwrap()
|
||||
});
|
||||
// A bad later batch must abort the request, including after
|
||||
// an earlier valid batch was encoded. No write handler runs.
|
||||
let mut batches = Vec::new();
|
||||
if range.is_some() {
|
||||
let empty_range_arrays: Vec<ArrayRef> = vec![
|
||||
Arc::new(LargeBinaryArray::from(vec![Some(b"first".as_slice())])),
|
||||
Arc::new(StringArray::from(vec![None::<&str>])),
|
||||
Arc::new(UInt64Array::new_null(1)),
|
||||
Arc::new(UInt64Array::new_null(1)),
|
||||
];
|
||||
let DataType::Struct(children) = schema.field(1).data_type().clone() else {
|
||||
unreachable!()
|
||||
};
|
||||
batches.push(Ok(RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int64Array::from(vec![0])),
|
||||
Arc::new(StructArray::new(children, empty_range_arrays, None)),
|
||||
],
|
||||
)
|
||||
.unwrap()));
|
||||
}
|
||||
batches.push(Ok(batch));
|
||||
let reader: Box<dyn RecordBatchReader + Send> =
|
||||
Box::new(RecordBatchIterator::new(batches, schema.clone()));
|
||||
let result = match operation {
|
||||
"merge" => {
|
||||
let mut merge = table.merge_insert(&["id"]);
|
||||
merge
|
||||
.when_matched_update_all(None)
|
||||
.when_not_matched_insert_all();
|
||||
merge.execute(reader).await.map(|_| ())
|
||||
}
|
||||
"overwrite" => table
|
||||
.add(reader)
|
||||
.mode(AddDataMode::Overwrite)
|
||||
.execute()
|
||||
.await
|
||||
.map(|_| ()),
|
||||
_ => table.add(reader).execute().await.map(|_| ()),
|
||||
};
|
||||
if !wide_target && range.is_some() && operation != "overwrite" {
|
||||
assert!(
|
||||
result
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("cannot store byte ranges")
|
||||
);
|
||||
assert!(
|
||||
receiver.try_recv().is_err(),
|
||||
"invalid stream must not reach a committed write"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
result.unwrap();
|
||||
let body = collect_body(receiver.recv().unwrap()).await;
|
||||
let reader = StreamReader::try_new(Cursor::new(body), None).unwrap();
|
||||
let expected = if operation == "overwrite" {
|
||||
schema.as_ref()
|
||||
} else {
|
||||
&target
|
||||
};
|
||||
assert_eq!(reader.schema().as_ref(), expected);
|
||||
for batch in reader {
|
||||
let batch = batch.unwrap();
|
||||
let image = batch.column(1).as_struct();
|
||||
assert!(
|
||||
image.column(0).as_binary::<i64>().value(0) == b"inline"
|
||||
|| image.column(0).as_binary::<i64>().value(0) == b"first"
|
||||
);
|
||||
if image.num_columns() == 4 && range.is_none() {
|
||||
assert!(image.column(2).is_null(0));
|
||||
assert!(image.column(3).is_null(0));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(true)]
|
||||
#[case(false)]
|
||||
@@ -4852,6 +5110,12 @@ mod tests {
|
||||
));
|
||||
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
if request.url().path().ends_with("/describe/") {
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version": 1, "schema": {"fields": []}}"#)
|
||||
.unwrap();
|
||||
}
|
||||
if request.url().path() == "/v1/table/my_table/merge_insert/" {
|
||||
assert_eq!(request.method(), "POST");
|
||||
|
||||
@@ -4908,6 +5172,12 @@ mod tests {
|
||||
));
|
||||
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
if request.url().path().ends_with("/describe/") {
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version": 1, "schema": {"fields": []}}"#)
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
|
||||
let params = request.url().query_pairs().collect::<HashMap<_, _>>();
|
||||
@@ -4993,6 +5263,12 @@ mod tests {
|
||||
let num_deleted_rows = if delete_unmatched { 3 } else { 0 };
|
||||
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
if request.url().path().ends_with("/describe/") {
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_response(&Schema::empty()))
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
|
||||
assert_eq!(request.headers()[CONTENT_TYPE], ARROW_STREAM_CONTENT_TYPE);
|
||||
@@ -5063,6 +5339,12 @@ mod tests {
|
||||
));
|
||||
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
if request.url().path().ends_with("/describe/") {
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version": 1, "schema": {"fields": []}}"#)
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
|
||||
|
||||
// One repeated `on` per column, in the order the caller gave them.
|
||||
@@ -5134,6 +5416,12 @@ mod tests {
|
||||
|
||||
// Default parameters
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
if request.url().path().ends_with("/describe/") {
|
||||
return http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"version": 1, "schema": {"fields": []}}"#)
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/");
|
||||
|
||||
@@ -14405,6 +14693,20 @@ mod tests {
|
||||
));
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => {
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(body["branch"], "exp");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_response(&Schema::new(vec![Field::new(
|
||||
"id",
|
||||
DataType::Int32,
|
||||
false,
|
||||
)])))
|
||||
.unwrap()
|
||||
}
|
||||
"/v1/table/my_table/branches/create/" => http::Response::builder()
|
||||
.status(200)
|
||||
.body("{}".to_string())
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow_schema::{DataType, Fields, Schema};
|
||||
use arrow_schema::{DataType, Field, Fields, Schema};
|
||||
use lance::dataset::WriteMode;
|
||||
use lance_arrow::FieldExt;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::data::scannable::Scannable;
|
||||
@@ -233,6 +235,7 @@ pub struct PreprocessingOutput {
|
||||
/// Fields can be in different orders, so match by name.
|
||||
///
|
||||
/// If a column exists in input but not in table, error (no extra columns allowed).
|
||||
/// Input field names must be unique within each struct.
|
||||
///
|
||||
/// If a column exists in table but not in input, that is okay - it may be filled with nulls.
|
||||
///
|
||||
@@ -244,7 +247,13 @@ fn validate_schema(input: &Schema, table: &Schema) -> Result<()> {
|
||||
}
|
||||
|
||||
fn validate_fields(input: &Fields, table: &Fields) -> Result<()> {
|
||||
let mut names = HashSet::with_capacity(input.len());
|
||||
for field in input {
|
||||
if !names.insert(field.name()) {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!("field '{}' is specified more than once", field.name()),
|
||||
});
|
||||
}
|
||||
match table.iter().find(|f| f.name() == field.name()) {
|
||||
None => {
|
||||
return Err(Error::InvalidInput {
|
||||
@@ -252,17 +261,33 @@ fn validate_fields(input: &Fields, table: &Fields) -> Result<()> {
|
||||
});
|
||||
}
|
||||
Some(table_field) => {
|
||||
if let (DataType::Struct(in_children), DataType::Struct(tbl_children)) =
|
||||
(field.data_type(), table_field.data_type())
|
||||
{
|
||||
validate_fields(in_children, tbl_children)?;
|
||||
}
|
||||
validate_field(field, table_field)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_field(input: &Field, table: &Field) -> Result<()> {
|
||||
// Blob coercion validates its children and permits dropping only
|
||||
// all-null range fields when targeting a legacy two-field layout.
|
||||
if table.is_blob_v2() {
|
||||
return Ok(());
|
||||
}
|
||||
match (input.data_type(), table.data_type()) {
|
||||
(DataType::Struct(input), DataType::Struct(table)) => validate_fields(input, table),
|
||||
(
|
||||
DataType::List(input) | DataType::LargeList(input) | DataType::FixedSizeList(input, _),
|
||||
DataType::List(table) | DataType::LargeList(table) | DataType::FixedSizeList(table, _),
|
||||
) => {
|
||||
// Preserve the item's extension metadata, and match nested struct
|
||||
// children by name rather than comparing collection item names.
|
||||
validate_field(input, table)
|
||||
}
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
@@ -299,6 +324,81 @@ mod tests {
|
||||
conn.create_table("test", batch).execute().await.unwrap()
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
fn test_validate_append_collection_fields(
|
||||
#[values("list", "large_list", "fixed_size_list")] source_kind: &str,
|
||||
#[values("list", "large_list", "fixed_size_list")] target_kind: &str,
|
||||
) {
|
||||
fn schema(kind: &str, item: Field) -> Schema {
|
||||
let item = Arc::new(item);
|
||||
let data_type = match kind {
|
||||
"list" => DataType::List(item),
|
||||
"large_list" => DataType::LargeList(item),
|
||||
"fixed_size_list" => DataType::FixedSizeList(item, 1),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
Schema::new(vec![Field::new("items", data_type, true)])
|
||||
}
|
||||
let image = crate::blob("image", true);
|
||||
let legacy_image = lance::blob::blob_field("image", true);
|
||||
let caption = Field::new("caption", DataType::Utf8, true);
|
||||
let input = schema(
|
||||
source_kind,
|
||||
Field::new(
|
||||
"source_item",
|
||||
DataType::Struct(vec![image.clone(), caption.clone()].into()),
|
||||
true,
|
||||
),
|
||||
);
|
||||
let target = schema(
|
||||
target_kind,
|
||||
Field::new(
|
||||
"target_item",
|
||||
DataType::Struct(vec![legacy_image.clone()].into()),
|
||||
true,
|
||||
),
|
||||
);
|
||||
let error = super::validate_schema(&input, &target)
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
assert!(
|
||||
error.contains("caption") && error.contains("does not exist"),
|
||||
"{error}"
|
||||
);
|
||||
|
||||
let target = schema(
|
||||
target_kind,
|
||||
Field::new(
|
||||
"target_item",
|
||||
DataType::Struct(vec![caption.clone(), legacy_image.clone()].into()),
|
||||
true,
|
||||
),
|
||||
);
|
||||
// Declared siblings can be reordered, but must not occur more than once.
|
||||
super::validate_schema(&input, &target).unwrap();
|
||||
let duplicate = schema(
|
||||
source_kind,
|
||||
Field::new(
|
||||
"source_item",
|
||||
DataType::Struct(vec![image.clone(), caption.clone(), caption].into()),
|
||||
true,
|
||||
),
|
||||
);
|
||||
let error = super::validate_schema(&duplicate, &target)
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
assert!(
|
||||
error.contains("caption") && error.contains("specified more than once"),
|
||||
"{error}"
|
||||
);
|
||||
// Item names need not match, and blob children still reach range coercion.
|
||||
super::validate_schema(
|
||||
&schema(source_kind, image.with_name("source_item")),
|
||||
&schema(target_kind, legacy_image.with_name("target_item")),
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn test_add_with_data<T>(data: T)
|
||||
where
|
||||
T: Scannable + 'static,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -33,9 +33,12 @@ pub fn cast_to_table_schema(
|
||||
return Ok(input);
|
||||
}
|
||||
|
||||
let exprs = build_field_exprs(input_schema.fields(), table_schema.fields(), &|idx| {
|
||||
Arc::new(Column::new(input_schema.field(idx).name(), idx)) as Arc<dyn PhysicalExpr>
|
||||
})?;
|
||||
let exprs = build_field_exprs(
|
||||
input_schema.fields(),
|
||||
table_schema.fields(),
|
||||
&|idx| Arc::new(Column::new(input_schema.field(idx).name(), idx)) as Arc<dyn PhysicalExpr>,
|
||||
&[],
|
||||
)?;
|
||||
|
||||
let exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = exprs
|
||||
.into_iter()
|
||||
@@ -52,10 +55,13 @@ pub fn cast_to_table_schema(
|
||||
/// For each table field that exists in the input, produce an expression that
|
||||
/// reads from the input and casts if needed. Fields in the table but not in the
|
||||
/// input are omitted (the storage layer handles missing columns).
|
||||
fn build_field_exprs(
|
||||
/// `ancestors` carries validity for blob range validation without nullifying
|
||||
/// non-nullable children hidden beneath a null parent.
|
||||
pub(super) fn build_field_exprs(
|
||||
input_fields: &Fields,
|
||||
table_fields: &Fields,
|
||||
get_input_expr: &dyn Fn(usize) -> Arc<dyn PhysicalExpr>,
|
||||
ancestors: &[Arc<dyn PhysicalExpr>],
|
||||
) -> Result<Vec<(Arc<dyn PhysicalExpr>, FieldRef)>> {
|
||||
let config = Arc::new(ConfigOptions::default());
|
||||
let mut result = Vec::new();
|
||||
@@ -109,13 +115,14 @@ fn build_field_exprs(
|
||||
input_field,
|
||||
table_field,
|
||||
&config,
|
||||
ancestors,
|
||||
)?);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Blob values in lists need the same coercion as top-level blob columns.
|
||||
if let Some(coerced) =
|
||||
coerce_blob_list_expr(input_expr.clone(), input_field, table_field, &config)?
|
||||
coerce_blob_list_expr(input_expr.clone(), input_field, table_field, ancestors)?
|
||||
{
|
||||
result.push(coerced);
|
||||
continue;
|
||||
@@ -138,30 +145,38 @@ fn build_field_exprs(
|
||||
(DataType::Struct(in_children), DataType::Struct(tbl_children))
|
||||
if in_children != tbl_children =>
|
||||
{
|
||||
let sub_exprs = build_field_exprs(in_children, tbl_children, &|child_idx| {
|
||||
let child_name = in_children[child_idx].name();
|
||||
Arc::new(ScalarFunctionExpr::new(
|
||||
&format!("get_field({child_name})"),
|
||||
get_field(),
|
||||
vec![
|
||||
input_expr.clone(),
|
||||
Arc::new(Literal::new(ScalarValue::from(child_name.as_str()))),
|
||||
],
|
||||
Arc::new(in_children[child_idx].as_ref().clone()),
|
||||
config.clone(),
|
||||
)) as Arc<dyn PhysicalExpr>
|
||||
})?;
|
||||
let mut child_ancestors = ancestors.to_vec();
|
||||
child_ancestors.push(input_expr.clone());
|
||||
let sub_exprs = build_field_exprs(
|
||||
in_children,
|
||||
tbl_children,
|
||||
&|child_idx| {
|
||||
let child_name = in_children[child_idx].name();
|
||||
Arc::new(ScalarFunctionExpr::new(
|
||||
&format!("get_field({child_name})"),
|
||||
get_field(),
|
||||
vec![
|
||||
input_expr.clone(),
|
||||
Arc::new(Literal::new(ScalarValue::from(child_name.as_str()))),
|
||||
],
|
||||
Arc::new(in_children[child_idx].as_ref().clone()),
|
||||
config.clone(),
|
||||
)) as Arc<dyn PhysicalExpr>
|
||||
},
|
||||
&child_ancestors,
|
||||
)?;
|
||||
|
||||
let output_struct_fields: Fields = sub_exprs
|
||||
.iter()
|
||||
.map(|(_, f)| f.clone())
|
||||
.collect::<Vec<_>>()
|
||||
.into();
|
||||
let output_field: FieldRef = Arc::new(Field::new(
|
||||
table_field.name(),
|
||||
DataType::Struct(output_struct_fields),
|
||||
table_field.is_nullable(),
|
||||
));
|
||||
let output_field: FieldRef = Arc::new(
|
||||
table_field
|
||||
.as_ref()
|
||||
.clone()
|
||||
.with_data_type(DataType::Struct(output_struct_fields)),
|
||||
);
|
||||
|
||||
// Build named_struct(lit("a"), expr_a, lit("b"), expr_b, ...)
|
||||
let mut ns_args: Vec<Arc<dyn PhysicalExpr>> = Vec::new();
|
||||
|
||||
@@ -254,7 +254,11 @@ fn canonicalize_merge_filter(filter: Option<MergeFilter>) -> Result<Option<Merge
|
||||
// The JSON projection iterates target fields, so reject unknown source fields
|
||||
// before it can erase them. Missing and reordered fields remain valid merge
|
||||
// inputs; type compatibility is still checked by the cast and Lance writer.
|
||||
fn validate_merge_source_fields(input: &Fields, target: &Fields, parent: &str) -> Result<()> {
|
||||
pub(crate) fn validate_merge_source_fields(
|
||||
input: &Fields,
|
||||
target: &Fields,
|
||||
parent: &str,
|
||||
) -> Result<()> {
|
||||
let mut names = HashSet::with_capacity(input.len());
|
||||
for field in input {
|
||||
let path = if parent.is_empty() {
|
||||
@@ -273,7 +277,9 @@ fn validate_merge_source_fields(input: &Fields, target: &Fields, parent: &str) -
|
||||
.ok_or_else(|| Error::InvalidInput {
|
||||
message: format!("merge source field '{path}' is not present in the table schema"),
|
||||
})?;
|
||||
validate_merge_source_type(field.data_type(), target_field.data_type(), &path)?;
|
||||
if !crate::blob::is_blob(target_field) {
|
||||
validate_merge_source_type(field.data_type(), target_field.data_type(), &path)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -296,7 +302,13 @@ fn validate_merge_source_type(input: &DataType, target: &DataType, path: &str) -
|
||||
| DataType::LargeListView(target),
|
||||
)
|
||||
| (DataType::Map(input, _), DataType::Map(target, _)) => {
|
||||
validate_merge_source_type(input.data_type(), target.data_type(), path)
|
||||
// Collection items carry extension metadata just like named fields.
|
||||
// Blob coercion, rather than struct validation, checks range narrowing.
|
||||
if crate::blob::is_blob(target) {
|
||||
Ok(())
|
||||
} else {
|
||||
validate_merge_source_type(input.data_type(), target.data_type(), path)
|
||||
}
|
||||
}
|
||||
(DataType::Dictionary(_, input), _) => validate_merge_source_type(input, target, path),
|
||||
(_, DataType::Dictionary(_, target)) => validate_merge_source_type(input, target, path),
|
||||
@@ -313,35 +325,18 @@ pub(crate) async fn execute_merge_insert(
|
||||
new_data: Box<dyn RecordBatchReader + Send>,
|
||||
) -> Result<MergeResult> {
|
||||
params.canonicalize_filters()?;
|
||||
let schema = table.schema().await?;
|
||||
super::computed_columns::ensure_no_function_bindings_for_mutation(
|
||||
table.schema().await?.as_ref(),
|
||||
schema.as_ref(),
|
||||
"merge_insert",
|
||||
)?;
|
||||
match lsm::lsm_dispatch_decision(table, ¶ms).await? {
|
||||
lsm::LsmDispatch::Lsm(plan) => {
|
||||
let future =
|
||||
lsm::execute_lsm_merge_insert(table, plan, params.validate_single_shard, new_data);
|
||||
return match params.timeout {
|
||||
Some(timeout) => match tokio::time::timeout(timeout, future).await {
|
||||
Ok(result) => result,
|
||||
Err(_) => Err(Error::Runtime {
|
||||
message: "merge insert timed out".to_string(),
|
||||
}),
|
||||
},
|
||||
None => future.await,
|
||||
};
|
||||
}
|
||||
lsm::LsmDispatch::Standard => {}
|
||||
}
|
||||
|
||||
let dataset = table.dataset.get().await?;
|
||||
let schema = arrow_schema::Schema::from(dataset.schema());
|
||||
// JSON source fields must carry the stored extension metadata just as on
|
||||
// append. Keep arrow.json text labelled until Lance encodes it as JSONB.
|
||||
let source = if schema
|
||||
.fields()
|
||||
.iter()
|
||||
.any(|field| lance_arrow::json::has_json_fields(field))
|
||||
// Normalize blob layouts and JSON extension metadata just as on append.
|
||||
// Keep arrow.json text labelled until Lance encodes it as JSONB.
|
||||
let source = if crate::blob::has_blob_columns(&schema)
|
||||
|| schema
|
||||
.fields()
|
||||
.iter()
|
||||
.any(|field| lance_arrow::json::has_json_fields(field))
|
||||
{
|
||||
validate_merge_source_fields(new_data.schema().fields(), schema.fields(), "")?;
|
||||
let plan = Arc::new(super::datafusion::scannable_exec::ScannableExec::new(
|
||||
@@ -356,6 +351,24 @@ pub(crate) async fn execute_merge_insert(
|
||||
} else {
|
||||
new_data.into_stream()
|
||||
};
|
||||
match lsm::lsm_dispatch_decision(table, ¶ms).await? {
|
||||
lsm::LsmDispatch::Lsm(plan) => {
|
||||
let future =
|
||||
lsm::execute_lsm_merge_insert(table, plan, params.validate_single_shard, source);
|
||||
return match params.timeout {
|
||||
Some(timeout) => match tokio::time::timeout(timeout, future).await {
|
||||
Ok(result) => result,
|
||||
Err(_) => Err(Error::Runtime {
|
||||
message: "merge insert timed out".to_string(),
|
||||
}),
|
||||
},
|
||||
None => future.await,
|
||||
};
|
||||
}
|
||||
lsm::LsmDispatch::Standard => {}
|
||||
}
|
||||
|
||||
let dataset = table.dataset.get().await?;
|
||||
let mut builder = LanceMergeInsertBuilder::try_new(dataset.clone(), params.on)?;
|
||||
match (
|
||||
params.when_matched_update_all,
|
||||
|
||||
@@ -24,8 +24,10 @@ use arrow_array::cast::AsArray;
|
||||
use arrow_array::types::{
|
||||
Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
|
||||
};
|
||||
use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, RecordBatchReader};
|
||||
use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch};
|
||||
use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef};
|
||||
use datafusion_physical_plan::SendableRecordBatchStream;
|
||||
use futures::TryStreamExt;
|
||||
use lance::Dataset;
|
||||
use lance::dataset::mem_wal::{
|
||||
DatasetMemWalExt, ShardWriter, ShardWriterConfig, evaluate_sharding_spec,
|
||||
@@ -707,7 +709,7 @@ pub(crate) async fn execute_lsm_merge_insert(
|
||||
table: &NativeTable,
|
||||
plan: LsmPlan,
|
||||
validate_single_shard: bool,
|
||||
new_data: Box<dyn RecordBatchReader + Send>,
|
||||
mut new_data: SendableRecordBatchStream,
|
||||
) -> Result<MergeResult> {
|
||||
let dataset = table.dataset.get().await?;
|
||||
let target_schema: SchemaRef = Arc::new(ArrowSchema::from(dataset.schema()));
|
||||
@@ -718,8 +720,7 @@ pub(crate) async fn execute_lsm_merge_insert(
|
||||
let mut batches: Vec<RecordBatch> = Vec::new();
|
||||
let mut total_rows: u64 = 0;
|
||||
|
||||
for batch in new_data {
|
||||
let batch = batch.map_err(|e| Error::Arrow { source: e })?;
|
||||
while let Some(batch) = new_data.try_next().await? {
|
||||
if batch.num_rows() == 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -1747,7 +1747,7 @@ mod tests {
|
||||
RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", arrow_schema::DataType::Int32, false),
|
||||
crate::blob("image", true),
|
||||
lance::blob::blob_field("image", true),
|
||||
])),
|
||||
vec![Arc::new(Int32Array::from(ids)), builder.finish().unwrap()],
|
||||
)
|
||||
@@ -1901,7 +1901,7 @@ mod tests {
|
||||
let mut blob_builder = lance::blob::BlobArrayBuilder::new(2);
|
||||
blob_builder.push_bytes(b"nested").unwrap();
|
||||
blob_builder.push_null().unwrap();
|
||||
let blob_field = crate::blob("image", true);
|
||||
let blob_field = lance::blob::blob_field("image", true);
|
||||
let metadata_fields = Fields::from(vec![blob_field.clone()]);
|
||||
let metadata = StructArray::new(
|
||||
metadata_fields.clone(),
|
||||
@@ -1968,7 +1968,7 @@ mod tests {
|
||||
blob_builder.push_bytes(b"a").unwrap();
|
||||
blob_builder.push_bytes(b"bb").unwrap();
|
||||
blob_builder.push_null().unwrap();
|
||||
let item = Arc::new(crate::blob("item", true));
|
||||
let item = Arc::new(lance::blob::blob_field("item", true));
|
||||
let images = ListArray::new(
|
||||
item.clone(),
|
||||
OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 3])),
|
||||
|
||||
@@ -10,12 +10,13 @@ use arrow_array::{
|
||||
use arrow_schema::{DataType, Field, Fields, Schema};
|
||||
use futures::TryStreamExt;
|
||||
use lance::Dataset;
|
||||
use lance::blob::blob_field as blob;
|
||||
use lance::dataset::WriteParams;
|
||||
use lance_file::version::{ConcreteFileVersion, LanceFileVersion};
|
||||
use lance_table::format::BasePath;
|
||||
use lancedb::{
|
||||
Connection, Error, Result, Table,
|
||||
blob::{BlobRangeRequest, blob},
|
||||
blob::BlobRangeRequest,
|
||||
connect, connect_namespace,
|
||||
database::listing::{
|
||||
ListingDatabaseOptions, NewTableConfig, OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
|
||||
@@ -1610,6 +1611,88 @@ async fn external_uri_struct_round_trips_with_flag() -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn external_ranges_round_trip_with_registered_base_and_opt_in() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
let payload = b"0123456789abcdef";
|
||||
let uri = write_payload_file_uri(tmp.path(), "payload.bin", payload);
|
||||
let base_uri = url::Url::from_directory_path(tmp.path())
|
||||
.unwrap()
|
||||
.to_string();
|
||||
let db = connect(tmp.path().join("db").to_str().unwrap())
|
||||
.execute()
|
||||
.await?;
|
||||
|
||||
for registered_base in [false, true] {
|
||||
let field = lancedb::blob("image", true);
|
||||
let DataType::Struct(children) = field.data_type().clone() else {
|
||||
unreachable!("blob field is a struct");
|
||||
};
|
||||
let schema = Arc::new(Schema::new(vec![field]));
|
||||
let name = if registered_base {
|
||||
"registered"
|
||||
} else {
|
||||
"opt_in"
|
||||
};
|
||||
let table = db
|
||||
.create_empty_table(name, schema.clone())
|
||||
.write_options(WriteOptions {
|
||||
lance_write_params: Some(WriteParams {
|
||||
initial_bases: registered_base.then(|| {
|
||||
vec![BasePath {
|
||||
id: 1,
|
||||
name: Some("external".to_string()),
|
||||
path: base_uri.clone(),
|
||||
is_dataset_root: false,
|
||||
}]
|
||||
}),
|
||||
..Default::default()
|
||||
}),
|
||||
})
|
||||
.execute()
|
||||
.await?;
|
||||
let values = StructArray::new(
|
||||
children,
|
||||
vec![
|
||||
Arc::new(LargeBinaryArray::from(vec![
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(b"inline".as_slice()),
|
||||
])),
|
||||
Arc::new(StringArray::from(vec![
|
||||
Some(uri.as_str()),
|
||||
Some(uri.as_str()),
|
||||
Some(uri.as_str()),
|
||||
None,
|
||||
])),
|
||||
Arc::new(UInt64Array::from(vec![Some(0), Some(4), None, None])),
|
||||
Arc::new(UInt64Array::from(vec![Some(6), Some(6), None, None])),
|
||||
],
|
||||
None,
|
||||
);
|
||||
table
|
||||
.add(RecordBatch::try_new(schema, vec![Arc::new(values)])?)
|
||||
.allow_external_blob_outside_bases(!registered_base)
|
||||
.execute()
|
||||
.await?;
|
||||
let reopened = db.open_table(name).execute().await?;
|
||||
let ids = collect_row_ids(&reopened).await?;
|
||||
let bytes = reopened.fetch_blobs("image", &ids).await?;
|
||||
assert_eq!(bytes.value(0), &payload[..6]);
|
||||
assert_eq!(bytes.value(1), &payload[4..10]);
|
||||
assert_eq!(bytes.value(2), payload);
|
||||
assert_eq!(bytes.value(3), b"inline");
|
||||
let files = reopened.fetch_blob_files("image", &ids).await?;
|
||||
assert_eq!(files[1].as_ref().unwrap().size(), 6);
|
||||
assert_eq!(
|
||||
files[1].as_ref().unwrap().read_range(1..3).await?,
|
||||
&payload[5..7]
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn external_uri_add_requires_opt_in() -> Result<()> {
|
||||
let tmp = tempdir().unwrap();
|
||||
@@ -1761,3 +1844,188 @@ async fn malformed_string_uri_is_rejected_at_write() -> Result<()> {
|
||||
assert_eq!(table.count_rows(None).await?, 0);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn layout_batch(wide: bool, id: i64, position: Option<u64>, size: Option<u64>) -> RecordBatch {
|
||||
let field = if wide {
|
||||
lancedb::blob("image", true)
|
||||
} else {
|
||||
blob("image", true)
|
||||
};
|
||||
let DataType::Struct(children) = field.data_type().clone() else {
|
||||
unreachable!()
|
||||
};
|
||||
let mut columns: Vec<arrow_array::ArrayRef> = vec![
|
||||
Arc::new(LargeBinaryArray::from(vec![Some(b"payload".as_slice())])),
|
||||
new_null_array(&DataType::Utf8, 1),
|
||||
];
|
||||
if wide {
|
||||
columns.push(Arc::new(UInt64Array::from(vec![position])));
|
||||
columns.push(Arc::new(UInt64Array::from(vec![size])));
|
||||
}
|
||||
RecordBatch::try_new(
|
||||
Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
field,
|
||||
])),
|
||||
vec![
|
||||
Arc::new(Int64Array::from(vec![id])),
|
||||
Arc::new(StructArray::new(children, columns, None)),
|
||||
],
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn layout_reader(batches: Vec<RecordBatch>) -> Box<dyn arrow_array::RecordBatchReader + Send> {
|
||||
let schema = batches[0].schema();
|
||||
Box::new(arrow_array::RecordBatchIterator::new(
|
||||
batches.into_iter().map(Ok),
|
||||
schema,
|
||||
))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn blob_layout_append_and_merge_normalize_both_directions() -> Result<()> {
|
||||
for wide in [false, true] {
|
||||
for merge in [false, true] {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
|
||||
let initial = layout_batch(wide, 1, None, None);
|
||||
let schema = initial.schema();
|
||||
let table = db.create_table("layouts", initial).execute().await?;
|
||||
let input = layout_batch(!wide, if merge { 1 } else { 2 }, None, None);
|
||||
if merge {
|
||||
let mut builder = table.merge_insert(&["id"]);
|
||||
builder
|
||||
.when_matched_update_all(None)
|
||||
.when_not_matched_insert_all();
|
||||
let result = builder
|
||||
.execute(layout_reader(vec![
|
||||
input,
|
||||
layout_batch(!wide, 2, None, None),
|
||||
]))
|
||||
.await?;
|
||||
assert_eq!(result.num_updated_rows, 1);
|
||||
assert_eq!(result.num_inserted_rows, 1);
|
||||
} else {
|
||||
table.add(input).execute().await?;
|
||||
}
|
||||
assert_eq!(table.schema().await?, schema);
|
||||
assert_eq!(table.count_rows(None).await?, 2);
|
||||
let ids = collect_row_ids(&table).await?;
|
||||
let values = table.fetch_blobs("image", &ids).await?;
|
||||
assert_eq!(values.value(0), b"payload");
|
||||
assert_eq!(values.value(1), b"payload");
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn blob_layout_rejects_non_null_ranges_atomically_in_later_batches() -> Result<()> {
|
||||
for merge in [false, true] {
|
||||
for (position, size) in [(Some(0), None), (None, Some(6)), (Some(0), Some(6))] {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
|
||||
let table = db
|
||||
.create_table("layouts", layout_batch(false, 1, None, None))
|
||||
.execute()
|
||||
.await?;
|
||||
let version = table.version().await?;
|
||||
// A successful first batch must not commit when a later batch has a range.
|
||||
let input = layout_reader(vec![
|
||||
layout_batch(true, 2, None, None),
|
||||
layout_batch(true, 3, position, size),
|
||||
]);
|
||||
let result = if merge {
|
||||
let mut builder = table.merge_insert(&["id"]);
|
||||
builder
|
||||
.when_matched_update_all(None)
|
||||
.when_not_matched_insert_all();
|
||||
builder.execute(input).await.map(|_| ())
|
||||
} else {
|
||||
table.add(input).execute().await.map(|_| ())
|
||||
};
|
||||
let error = result.unwrap_err().to_string();
|
||||
assert!(error.contains("cannot store byte ranges"), "{error}");
|
||||
assert_eq!(table.version().await?, version);
|
||||
assert_eq!(table.count_rows(None).await?, 1);
|
||||
let ids = collect_row_ids(&table).await?;
|
||||
assert_eq!(table.fetch_blobs("image", &ids).await?.value(0), b"payload");
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn blob_layout_overwrite_replaces_schema_in_both_directions() -> Result<()> {
|
||||
for wide in [false, true] {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
|
||||
let table = db
|
||||
.create_table("layouts", layout_batch(wide, 1, None, None))
|
||||
.execute()
|
||||
.await?;
|
||||
let input = layout_batch(!wide, 2, None, None);
|
||||
let expected_schema = input.schema();
|
||||
table
|
||||
.add(input)
|
||||
.mode(AddDataMode::Overwrite)
|
||||
.execute()
|
||||
.await?;
|
||||
assert_eq!(table.schema().await?, expected_schema);
|
||||
assert_eq!(table.count_rows(None).await?, 1);
|
||||
let ids = collect_row_ids(&table).await?;
|
||||
assert_eq!(table.fetch_blobs("image", &ids).await?.value(0), b"payload");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn blob_layout_memwal_merge_normalizes_before_writing() -> Result<()> {
|
||||
use lancedb::{query::Select, table::MemWalWriteSpec};
|
||||
for wide in [false, true] {
|
||||
let tmp = tempdir().unwrap();
|
||||
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
|
||||
let table = db
|
||||
.create_table("layouts", layout_batch(wide, 1, None, None))
|
||||
.execute()
|
||||
.await?;
|
||||
table.set_unenforced_primary_key(["id"]).await?;
|
||||
table.enable_memwal(MemWalWriteSpec::unsharded()).await?;
|
||||
let mut merge = table.merge_insert(&["id"]);
|
||||
merge
|
||||
.when_matched_update_all(None)
|
||||
.when_not_matched_insert_all();
|
||||
merge
|
||||
.clone()
|
||||
.execute(layout_reader(vec![layout_batch(!wide, 2, None, None)]))
|
||||
.await?;
|
||||
let batches = table
|
||||
.query()
|
||||
.select(Select::columns(&["id"]))
|
||||
.execute()
|
||||
.await?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
|
||||
if !wide {
|
||||
let error = merge
|
||||
.execute(layout_reader(vec![
|
||||
layout_batch(true, 3, None, None),
|
||||
layout_batch(true, 4, Some(0), Some(6)),
|
||||
]))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.to_string().contains("cannot store byte ranges"));
|
||||
let batches = table
|
||||
.query()
|
||||
.select(Select::columns(&["id"]))
|
||||
.execute()
|
||||
.await?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -410,6 +410,8 @@ async fn dedup_query_materializes_blob_payloads_and_null_vectors() -> anyhow::Re
|
||||
(0..6).map(|id| vec![id; 4096]),
|
||||
)),
|
||||
Arc::new(StringArray::from(vec![None::<&str>; 6])),
|
||||
Arc::new(UInt64Array::new_null(6)),
|
||||
Arc::new(UInt64Array::new_null(6)),
|
||||
],
|
||||
None,
|
||||
);
|
||||
|
||||
Reference in New Issue
Block a user