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>

python/python/tests/test_elastic_dataloader.py

@@ -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)
 # ---------------------------------------------------------------------------

python/python/tests/test_permutation.py

@@ -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)}))

rust/lancedb/src/dataloader/permutation/builder.rs

@@ -3,6 +3,7 @@
 
 use std::{collections::HashMap, sync::Arc};
 
+use arrow_array::RecordBatch;
 use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_execution::{disk_manager::DiskManagerBuilder, runtime_env::RuntimeEnvBuilder};
 use datafusion_expr::col;
@@ -220,6 +221,12 @@ impl PermutationBuilder {
         Ok(Box::pin(SimpleRecordBatchStream { schema, stream }))
     }
 
+    /// Replaces the stream's schema metadata with the permutation config.
+    ///
+    /// The scan of the base table can carry the base table's schema metadata (e.g. Hugging
+    /// Face's `huggingface` key).  That metadata describes columns the permutation table does
+    /// not have, and its keys could collide with the config keys, so it is dropped rather
+    /// than merged.
     fn add_config_metadata(
         data: SendableRecordBatchStream,
         metadata: HashMap<String, String>,
@@ -227,7 +234,13 @@ impl PermutationBuilder {
         let schema = data.schema().as_ref().clone().with_metadata(metadata);
         let schema = Arc::new(schema);
         let schema_clone = schema.clone();
-        let stream = data.map_ok(move |batch| batch.with_schema(schema.clone()).unwrap());
+        // `RecordBatch::with_schema` only accepts a superset of the batch's schema, which a
+        // schema with replaced metadata is not, so rebuild the batch around the new schema.
+        let stream = data.and_then(move |batch| {
+            std::future::ready(
+                RecordBatch::try_new(schema.clone(), batch.columns().to_vec()).map_err(Error::from),
+            )
+        });
         Ok(Box::pin(SimpleRecordBatchStream {
             schema: schema_clone,
             stream,
@@ -776,4 +789,69 @@ mod tests {
             "expected the pre-check to refuse the table, got: {err}"
         );
     }
+
+    /// Base table schema metadata (e.g. Hugging Face's `huggingface` key) describes the base
+    /// table's columns, so the permutation table drops it and keeps only its own config.
+    #[tokio::test]
+    async fn test_permutation_over_table_with_schema_metadata() {
+        use arrow_array::{Int64Array, RecordBatchIterator};
+        use arrow_schema::{DataType, Field, Schema};
+
+        let temp_dir = tempfile::tempdir().unwrap();
+        let db = connect(temp_dir.path().to_str().unwrap())
+            .execute()
+            .await
+            .unwrap();
+        let schema = Arc::new(
+            Schema::new(vec![Field::new("x", DataType::Int64, false)]).with_metadata(
+                HashMap::from([
+                    ("huggingface".to_string(), "{}".to_string()),
+                    // Collides with a config key; the config value must win.
+                    (SPLIT_NAMES_CONFIG_KEY.to_string(), "not json".to_string()),
+                ]),
+            ),
+        );
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int64Array::from_iter_values(0..100))],
+        )
+        .unwrap();
+        let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
+            Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
+        let table = db.create_table("tbl", reader).execute().await.unwrap();
+        assert!(
+            table
+                .schema()
+                .await
+                .unwrap()
+                .metadata
+                .contains_key("huggingface")
+        );
+
+        let permutation_table = PermutationBuilder::new(table)
+            .with_split_strategy(
+                SplitStrategy::Random {
+                    seed: Some(1),
+                    sizes: SplitSizes::Fixed(2),
+                    clump_size: None,
+                },
+                Some(vec!["train".to_string(), "test".to_string()]),
+            )
+            .with_shuffle_strategy(ShuffleStrategy::Random {
+                seed: Some(1),
+                clump_size: None,
+            })
+            .build()
+            .await
+            .unwrap();
+
+        assert_eq!(permutation_table.count_rows(None).await.unwrap(), 100);
+        let metadata = permutation_table.schema().await.unwrap().metadata.clone();
+        assert!(!metadata.contains_key("huggingface"));
+        assert_eq!(
+            metadata.get(SPLIT_NAMES_CONFIG_KEY).map(String::as_str),
+            Some(r#"["train","test"]"#)
+        );
+        assert!(metadata.contains_key(BASE_VERSION_CONFIG_KEY));
+    }
 }

rust/lancedb/src/dataloader/permutation/split.rs

@@ -214,9 +214,9 @@ impl Splitter {
             }
             arrays.push(Arc::new(split_ids));
 
-            std::future::ready(Some(Ok(
-                RecordBatch::try_new(new_schema.clone(), arrays).unwrap()
-            )))
+            std::future::ready(Some(
+                RecordBatch::try_new(new_schema.clone(), arrays).map_err(Error::from),
+            ))
         });
 
         Ok(Box::pin(SimpleRecordBatchStream::new(
@@ -289,23 +289,23 @@ impl Splitter {
             .collect::<Vec<_>>();
 
         let new_schema_clone = new_schema.clone();
-        let stream = source.map_ok(move |batch| {
+        let stream = source.and_then(move |batch| {
             let split_ids = Self::hash_split_id(&batch, &thresholds, total_weight);
 
-            if split_ids.null_count() > 0 {
-                let is_valid = split_ids.nulls().unwrap().inner();
-                let is_valid_mask = BooleanArray::new(is_valid.clone(), None);
-                let split_ids = arrow::compute::filter(&split_ids, &is_valid_mask).unwrap();
-                let row_ids = batch.column(row_id_index);
-                let row_ids = arrow::compute::filter(row_ids.as_ref(), &is_valid_mask).unwrap();
-                RecordBatch::try_new(new_schema.clone(), vec![row_ids, split_ids]).unwrap()
+            let batch = if let Some(nulls) = split_ids.nulls().filter(|n| n.null_count() > 0) {
+                let is_valid_mask = BooleanArray::new(nulls.inner().clone(), None);
+                arrow::compute::filter(&split_ids, &is_valid_mask).and_then(|split_ids| {
+                    let row_ids = batch.column(row_id_index);
+                    let row_ids = arrow::compute::filter(row_ids.as_ref(), &is_valid_mask)?;
+                    RecordBatch::try_new(new_schema.clone(), vec![row_ids, split_ids])
+                })
             } else {
                 RecordBatch::try_new(
                     new_schema.clone(),
                     vec![batch.column(row_id_index).clone(), Arc::new(split_ids)],
                 )
-                .unwrap()
-            }
+            };
+            std::future::ready(batch.map_err(Error::from))
         });
 
         Ok(Box::pin(SimpleRecordBatchStream::new(