fix: support blob v2 in fragment update_columns - #8344
fix: support blob v2 in fragment update_columns#8344lance-gatefixer[bot] wants to merge 5 commits into
Conversation
There was a problem hiding this comment.
❌ Gate recommendation: request changes.
The top-level descriptor-to-logical mechanism fixes the reported inline/null/empty case, but selected-field traversal and fallback rewriting still reject two valid blob-v2 shapes. A complete revision should apply descriptor handling recursively and preserve trusted existing absolute external references without implicitly allowing new unregistered right-hand values.
| let descriptor_blob_ids = read_columns | ||
| .iter() | ||
| .filter_map(|column| self.schema().field(column)) | ||
| .filter(|field| field.is_blob_v2()) |
There was a problem hiding this comment.
Nested blob-v2 fields still bypass this detector: schema.field(column) resolves the selected top-level struct, whose child may be blob-v2 even though the parent is not. Updating such a struct therefore reaches the original decoder mismatch. Drive both descriptor projection and descriptor-to-logical conversion from a recursive search of each selected field tree.
Reproducer
I ran this against 38df2a140e2c2624d0e25fa73ecf18a870544af9:
import tempfile
from pathlib import Path
import lance
import pyarrow as pa
def info(names, payloads):
fields = [pa.field("name", pa.string()), lance.blob_field("blob")]
return pa.StructArray.from_arrays(
[pa.array(names), lance.blob_array(payloads)], fields=fields
)
with tempfile.TemporaryDirectory() as root:
uri = Path(root) / "nested.lance"
ds = lance.write_dataset(
pa.table({"id": [1, 2], "info": info(["a", "b"], [b"one", b"two"])}),
uri, data_storage_version="2.2",
)
ds.get_fragment(0).update_columns(
pa.table({"id": [2], "info": info(["B"], [b"NEW"])}),
left_on="id",
)Expected: the update returns fragment metadata. Observed: ValueError: ... more fields in the schema than provided column indices / infos.
There was a problem hiding this comment.
Addressed in 1014ee2: selected blob-v2 fields and descriptor-to-logical conversion now traverse nested structs and lists recursively, with regression coverage for the nested-struct reproducer.
| })?; | ||
| while let Some(batch) = updater.next().await? { | ||
| let batch = if has_blob_v2 { | ||
| crate::dataset::optimize::transform_blob_v2_batch( |
There was a problem hiding this comment.
This conversion exposes a stored base_id == 0 fallback descriptor as its absolute URI, but the updater opens its writer with default parameters, where allow_external_blob_outside_bases is false. An update then fails merely because an untouched row contains a valid absolute external blob. Preserve trusted existing references separately from incoming values—for example, carry their descriptors through fallback or allowlist only references already present—so fixing this does not silently opt new right-hand URIs into a broader write policy.
Reproducer
I ran this against 38df2a140e2c2624d0e25fa73ecf18a870544af9:
import tempfile
from pathlib import Path
import lance
import pyarrow as pa
with tempfile.TemporaryDirectory() as root, tempfile.TemporaryDirectory() as outside:
uri = Path(root) / "external.lance"
external = Path(outside) / "payload.bin"
external.write_bytes(b"outside")
ds = lance.write_dataset(
pa.table({"id": [1, 2], "payload": lance.blob_array([external.as_uri(), b"two"])}),
uri,
data_storage_version="2.2",
allow_external_blob_outside_bases=True,
)
ds.get_fragment(0).update_columns(
pa.table({"id": [2], "payload": lance.blob_array([b"NEW"])}),
left_on="id",
)Expected: row 2 updates while row 1 keeps its existing external reference. Observed: Invalid user input: External blob URI ... is outside registered external bases.
There was a problem hiding this comment.
Addressed in 1014ee2: unmatched absolute external references are preserved, while matched right-hand URIs are validated against registered bases before the update writer accepts trusted fallback references. Regression coverage verifies both preservation and rejection of a new unregistered URI.
There was a problem hiding this comment.
❌ Gate recommendation: request changes.
The recursive descriptor handling and external-reference boundary are now correct, but the shared rewrite path regresses bounded memory for both fragment updates and compaction. Keep the field-ID-aware recursion while streaming each managed payload directly into the final logical builder, so a rewrite never retains two whole payload batches.
| .map(|index| row_addrs[*index]) | ||
| .collect::<Vec<_>>(); | ||
| Some( | ||
| super::blob::materialize_blob_v2_descriptors( |
There was a problem hiding this comment.
materialize_blob_v2_descriptors defeats the lazy-memory contract here: it collects every managed blob, copies the whole batch into a LargeBinaryArray, and the loop below copies that array again into data_builder. The previous path read one BlobFile at a time, so peak payload residency changes from the output plus one value to roughly two complete batches. Because this helper is shared with compaction and updater batches default to 8,192 rows, multi-megabyte blobs can now turn an otherwise bounded rewrite into OOM.
Keep the recursive field-ID resolution, but expose per-descriptor BlobFile/Bytes values and append them directly to the final logical builder, avoiding the intermediate batch-wide array.
Reproducer
I ran these as separate processes at 1014ee202a2754844a6b3bb74b4e77e8fe8db1ef so setup memory was excluded from the measurement:
export PR8344_BENCH_DIR=$(mktemp -d /home/agent/tmp/pr8344-rss.XXXXXX)
uv run python - <<"PY"
import os
from pathlib import Path
import lance
import pyarrow as pa
uri = Path(os.environ["PR8344_BENCH_DIR"]) / "data.lance"
payload = b"x" * (8 * 1024 * 1024)
lance.write_dataset(
pa.table({"id": pa.array(range(32), type=pa.int64()), "payload": lance.blob_array([payload] * 32)}),
uri,
data_storage_version="2.2",
)
PY
uv run python - <<"PY"
import os
from pathlib import Path
import resource
import lance
import pyarrow as pa
uri = Path(os.environ["PR8344_BENCH_DIR"]) / "data.lance"
ds = lance.dataset(uri)
fragment, fields = ds.get_fragment(0).update_columns(
pa.table({"id": pa.array([0], type=pa.int64()), "payload": lance.blob_array([b"updated"])}),
left_on="id",
)
assert fields and fragment.physical_rows == 32
print(f"payload_mib={32 * 8} max_rss_kib={resource.getrusage(resource.RUSAGE_SELF).ru_maxrss}")
PYObserved on Linux: payload_mib=256 max_rss_kib=717072. The 256 MiB intermediate remains live while the final 256 MiB buffer is allocated and copied; larger batches scale linearly.
There was a problem hiding this comment.
Addressed in 5398dd3: descriptor conversion now retains only lazy, field-ID-aware BlobFile handles and reads one managed payload immediately before appending it to a pre-sized final logical buffer, eliminating the batch-wide intermediate LargeBinaryArray.
There was a problem hiding this comment.
❌ Gate recommendation: request changes.
The descriptor-aware update path is the right direction and the latest change removes the duplicate payload batch, but the recursive rewrite still lacks a single field-aware contract between its declared schema and array traversal. Build one recursive plan for both, with passthrough nodes for branches without blob-v2 descendants; that keeps the stream schema faithful and preserves bounded work.
|
|
||
| let new_struct = descriptor_to_logical_blob_array( | ||
| }; | ||
| let (new_column, new_field) = transform_blob_v2_array( |
There was a problem hiding this comment.
This now recursively transforms nested blob leaves in emitted batches, but transformed_schema below still changes a field only when the top-level Lance field is itself blob-v2. When Arrow JSON makes SchemaAdapter derive the write schema from that declared stream schema, the nested leaf remains descriptor-shaped and compaction fails.
Build the declared schema from the same recursive field plan used here so the batch and stream schemas cannot diverge.
Reproducer
I ran this at 5398dd3ca6f5812b18756f107e1136aafecc7851:
import json
import lance
import pyarrow as pa
def test_nested_blob_and_json_compaction(tmp_path):
uri = tmp_path / "nested-blob-json.lance"
fields = [lance.blob_field("blob"), pa.field("meta", pa.json_())]
schema = pa.schema([
pa.field("id", pa.int64()),
pa.field("info", pa.struct(fields)),
])
for index, row_id in enumerate([1, 2]):
info = pa.StructArray.from_arrays(
[
lance.blob_array([f"blob-{row_id}".encode()]),
pa.array([json.dumps({"row": row_id})], type=pa.json_()),
],
fields=fields,
)
lance.write_dataset(
pa.Table.from_arrays([pa.array([row_id]), info], schema=schema),
uri,
mode="create" if index == 0 else "append",
data_storage_version="2.2",
)
lance.dataset(uri).optimize.compact_files(num_threads=1)UV_NO_SYNC=1 UV_CACHE_DIR=/home/agent/tmp/pr8344-uv-cache \
PYTHONPATH=/home/agent/tmp/pr8344-impl/python/python \
uv run --project python pytest -p no:cacheprovider \
python/python/tests/test_gate_pr8344.py::test_nested_blob_and_json_compaction -qExpected: compaction completes. Observed: OSError: Invalid user input: Blob v2 field 'blob' has descriptor layout; expected logical or prepared layout.
There was a problem hiding this comment.
Addressed in 3673034: compaction now uses one recursive rewrite plan both to declare the stream schema and to transform every emitted batch. The nested blob-v2 plus Arrow JSON regression now compacts and reads both fields successfully.
| let values_len = values_end - values_start; | ||
| let mut normalized_offsets = Vec::with_capacity(list_array.len() + 1); | ||
| normalized_offsets.push(O::usize_as(0)); | ||
| let mut child_row_addrs = Vec::with_capacity(values_len); |
There was a problem hiding this comment.
This allocates one u64 row address per list element before checking whether the child tree contains a blob. Because transform_blob_v2_batch now invokes this transformer for every dataset field whenever any blob-v2 field exists, an unrelated one-byte blob makes large non-blob lists consume memory proportional to their element count and can turn compaction into an OOM.
Short-circuit each field and child subtree that has no blob-v2 descendant before rebuilding structs/lists or allocating child addresses. The same field-aware plan can also keep external-reference validation from walking unrelated branches.
Reproducer
I ran separate compaction processes at 5398dd3ca6f5812b18756f107e1136aafecc7851 over two fragments, each containing one list of five million int32 values, with batch_size=1. The only difference was an unrelated one-byte blob column:
import resource
import sys
from pathlib import Path
import lance
import numpy as np
import pyarrow as pa
mode, uri_arg, blob_arg = sys.argv[1:]
uri = Path(uri_arg)
with_blob = blob_arg == "blob"
def make_batch(row_id):
values = pa.array(np.arange(5_000_000, dtype=np.int32))
items = pa.ListArray.from_arrays(pa.array([0, len(values)]), values)
columns = {"id": pa.array([row_id]), "items": items}
if with_blob:
columns["payload"] = lance.blob_array([b"x"])
return pa.table(columns)
if mode == "create":
for row_id in range(2):
lance.write_dataset(
make_batch(row_id), uri,
mode="create" if row_id == 0 else "append",
data_storage_version="2.2",
)
else:
dataset = lance.dataset(uri)
dataset.optimize.compact_files(batch_size=1, num_threads=1)
print(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)Observed Linux ru_maxrss: 414824 KiB without the blob and 468708 KiB with it, an extra 53884 KiB for a field whose list subtree contains no blob.
There was a problem hiding this comment.
Addressed in 3673034: fields and child subtrees without blob-v2 descendants are now passthrough plan nodes, so unrelated lists are not rebuilt and receive no per-element row-address allocation. External-reference validation short-circuits the same unrelated branches.
There was a problem hiding this comment.
❌ Gate recommendation: request changes.
The shared recursive plan fixes the prior compaction schema and unrelated-list memory failures, but merge-insert still has a separate top-level-only advertised schema. Reuse the recursive plan and its output schema in merge-insert, as compaction now does, so partial-schema updates preserve nested blob-v2 fields without a descriptor/logical mismatch.
| batch: RecordBatch, | ||
| keep_row_addr: bool, | ||
| ) -> Result<RecordBatch> { | ||
| let plan = BlobV2BatchRewritePlan::try_new(schema, batch.schema().as_ref(), keep_row_addr)?; |
There was a problem hiding this comment.
This per-batch plan recursively converts nested blob descriptors to logical arrays, but merge-insert still advertises descriptor_to_logical_blob_schema, which converts only top-level blob fields. In a partial-schema update where info: struct<blob> is synthesized from the target, this returns a logical info array under a descriptor-shaped stream schema; create_filtered_batch then rejects the column and the merge cannot preserve the omitted field.
Reuse one BlobV2BatchRewritePlan, including its output_schema, in merge-insert just as compaction does, and add nested struct/list partial-schema merge coverage.
Reproducer
I ran this against 367303436de6ccb1b5a10eec10692fb817fbaf8a:
def test_partial_merge_insert_preserves_nested_blob(tmp_path):
uri = tmp_path / "partial-nested-blob-merge-insert.lance"
blob_field = lance.blob_field("blob")
info = pa.StructArray.from_arrays(
[lance.blob_array([b"one", b"two"])],
fields=[blob_field],
)
dataset = lance.write_dataset(
pa.table(
{
"id": pa.array([1, 2]),
"info": info,
"other": pa.array([10, 20]),
}
),
uri,
data_storage_version="2.2",
)
source = pa.table({"id": pa.array([2]), "other": pa.array([200])})
dataset.merge_insert("id").when_matched_update_all().execute(source)
result = lance.dataset(uri).to_table(blob_handling="all_binary").sort_by("id")
assert result["other"].to_pylist() == [10, 200]
assert result["info"].combine_chunks().field("blob").to_pylist() == [
b"one",
b"two",
]uv run --frozen pytest -p no:cacheprovider python/tests/test_gate_pr8344.py::test_partial_merge_insert_preserves_nested_blob -qExpected: the merge succeeds, updates other, and preserves the nested blobs. Observed: Invalid argument error: column types must match schema types, with a descriptor-shaped Struct expected and the logical blob Struct found at column index 1.
There was a problem hiding this comment.
Addressed in 942bfa1: merge-insert now reuses one recursive BlobV2BatchRewritePlan for both its advertised stream schema and per-batch transformation. Parameterized partial-schema regressions verify preservation of nested struct and list blob-v2 fields.
There was a problem hiding this comment.
✅ Gate recommendation: approve.
The latest revision removes the duplicate merge-insert schema converter and reuses one recursive blob rewrite plan for both the advertised stream schema and every emitted batch. Partial-schema struct and list regressions now preserve nested blob-v2 fields, while the previously verified compaction, external-reference, and bounded-memory behavior remains intact.
Summary
Root cause
FileFragment::update_columnsopened the logical blob-v2 struct schema directly against a blob encoded as one atomic physical column. The decoder therefore received more logical fields than projected physical column metadata. Existing fallback rows also need conversion from stored descriptors to the logical writer representation before they can be interleaved with incoming blob values.Validation
cargo fmt --all -- --checkcargo clippy --all --tests --benches -- -D warningscd python && make buildcd python && uv run make lintcd python && uv run pytest python/tests/test_fragment.py -k 'fragment_update_columns' -q(9 passed)Fixes #8336