mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-22 03:32:28 +00:00
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:
co-authored by
Claude Opus 5.5
parent
9f616e6218
commit
1772f4ef83
@@ -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)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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)}))
|
||||
|
||||
Reference in New Issue
Block a user