fix: don't panic building permutations over tables with schema metadata (#4602)

## Problem

`permutation_builder(table).execute()` and `StreamingDataset(table,
...)` (which builds a permutation at construction) panic with
`RustPanic` on any table whose Arrow schema has schema-level metadata:

```python
schema = pa.schema([pa.field("x", pa.int64())], metadata={b"huggingface": b"{}"})
t = lancedb.connect(tmp).create_table("t", pa.table({"x": list(range(100))}, schema=schema))
permutation_builder(t).split_random(fixed=2, seed=1).execute()   # RustPanic
```

Every Hugging Face Lance dataset triggers this, e.g.
`hf://datasets/lance-format/lerobot-xvla-soft-fold/data` table `frames`
carries `source-column-name-map`.

## Root cause

`PermutationBuilder::add_config_metadata` replaced the stream schema's
metadata with the permutation config and then called
`batch.with_schema(schema).unwrap()`. The base-table scan's batches
carry the base schema metadata, and `RecordBatch::with_schema` requires
the target to be a superset of the current schema, metadata included. It
returned `SchemaError("target schema is not superset of current schema
...")`, and the `unwrap` turned that into a panic on a tokio worker.

## Fix

- The permutation table deliberately does **not** inherit base-table
schema metadata. That metadata describes columns the permutation table
(`row_id`, `split_id`) does not have, and its keys can collide with
config keys such as `split_names`. A merge would let a base table's
`split_names` leak into, or override, the permutation's split names.
Batches are rebuilt with `RecordBatch::try_new(schema, columns)`.
- Errors are propagated (`map_ok` → `and_then`) instead of `unwrap`.
- The same panic-on-worker pattern in the splitter's streaming closures
(`split.rs`: `try_new`/`filter(...).unwrap()`) now propagates errors
too. Nothing else in `dataloader/` used `with_schema(...).unwrap()`.

## Tests

- Rust: `test_permutation_over_table_with_schema_metadata` in
`builder.rs` builds a random split + shuffled permutation over a table
with `huggingface` and a colliding `split_names` key. It panics with the
original `SchemaError` on the old code and passes now.
- Python: `test_permutation_over_table_with_schema_metadata` in
`test_permutation.py` covers the repro and reads both splits through
`Permutations`/`Permutation.from_tables`.
`test_streaming_dataset_over_table_with_schema_metadata` in
`test_elastic_dataloader.py` covers `StreamingDataset` construction and
iteration.
- `cargo fmt`, `cargo clippy --features remote --tests --examples`, the
dataloader Rust tests (42), `test_permutation.py` (72), and
`test_elastic_dataloader.py` all pass.

`test_torch.py::test_remote_permutation_dataloader_multiprocessing`
fails locally with `take query result did not include ordering column
'__lancedb_take_row_offset'`. It uses `Permutation.identity`, which does
not go through `PermutationBuilder`, so this change does not cause it.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Lei Xu
2026-10-20 09:16:18 -07:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 9f616e6218
commit 1772f4ef83
4 changed files with 137 additions and 14 deletions
@@ -3259,6 +3259,22 @@ def test_doc_example_checkpoint(lance_table):
)
def test_streaming_dataset_over_table_with_schema_metadata(tmp_path):
"""StreamingDataset builds its permutation at construction; base table schema
metadata (e.g. Hugging Face's `source-column-name-map`) must not break it."""
schema = pa.schema(
[pa.field("id", pa.int64())],
metadata={b"source-column-name-map": b"{}", b"huggingface": b"{}"},
)
table = lancedb.connect(tmp_path).create_table(
"data", pa.table({"id": list(range(NUM_ROWS))}, schema=schema)
)
ds = StreamingDataset(table, num_splits=NUM_SPLITS, shuffle_seed=SHUFFLE_SEED)
assert sorted(s["id"] for s in ds) == list(range(NUM_ROWS))
# ---------------------------------------------------------------------------
# Remote tables (LanceDB Cloud / Enterprise)
# ---------------------------------------------------------------------------
+29
View File
@@ -115,6 +115,35 @@ def test_split_random_fixed(mem_db):
assert split_ids.count(split_id) == 25
def test_permutation_over_table_with_schema_metadata(tmp_path):
"""Schema-level metadata on the base table must not break the build.
Hugging Face datasets carry keys like `huggingface` or `source-column-name-map`.
"""
schema = pa.schema(
[pa.field("x", pa.int64())],
# `split_names` collides with a permutation config key; the config must win.
metadata={b"huggingface": b"{}", b"split_names": b'["base"]'},
)
tbl = connect(tmp_path).create_table(
"t", pa.table({"x": list(range(100))}, schema=schema)
)
permutation_tbl = (
permutation_builder(tbl)
.split_random(fixed=2, seed=1, split_names=["train", "test"])
.execute()
)
permutations = Permutations(tbl, permutation_tbl)
assert permutations.split_names == ["train", "test"]
train = Permutation.from_tables(tbl, permutation_tbl, "train")
test = Permutation.from_tables(tbl, permutation_tbl, "test")
assert train.num_rows + test.num_rows == 100
train_x = [row["x"] for batch in train.iter(16) for row in batch]
test_x = [row["x"] for batch in test.iter(16) for row in batch]
assert sorted(train_x + test_x) == list(range(100))
def test_split_random_with_seed(mem_db):
"""Test that seeded random splits are reproducible."""
tbl = mem_db.create_table("test_table", pa.table({"x": range(50), "y": range(50)}))