fix!: normalize blob layouts and support external byte ranges (#4589)

This commit is contained in:
lancedb-gatefixer[bot]
2026-10-19 16:44:00 -07:00
committed by GitHub
parent 6507615a6f
commit 4a03859b2a
20 changed files with 2837 additions and 131 deletions
+9
View File
@@ -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.
+4
View File
@@ -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
+187 -1
View File
@@ -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([
+135
View File
@@ -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) => {
+187
View File
@@ -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
View File
@@ -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
View File
@@ -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;
}
+24 -12
View File
@@ -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".
+146
View File
@@ -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"
+83
View File
@@ -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",
+58 -5
View File
@@ -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]
+302
View File
@@ -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())
+106 -6
View File
@@ -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
+38 -23
View File
@@ -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();
+42 -29
View File
@@ -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, &params).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, &params).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,
+5 -4
View File
@@ -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;
}
+3 -3
View File
@@ -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])),
+269 -1
View File
@@ -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(())
}
+2
View File
@@ -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,
);