diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index cc6874e6c..9d18752f1 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/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) # --------------------------------------------------------------------------- diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index 142fc84f1..5e8bb1d76 100644 --- a/python/python/tests/test_permutation.py +++ b/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)})) diff --git a/rust/lancedb/src/dataloader/permutation/builder.rs b/rust/lancedb/src/dataloader/permutation/builder.rs index e08f3eddc..3837d9084 100644 --- a/rust/lancedb/src/dataloader/permutation/builder.rs +++ b/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, @@ -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 = + 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)); + } } diff --git a/rust/lancedb/src/dataloader/permutation/split.rs b/rust/lancedb/src/dataloader/permutation/split.rs index b33544100..cfa8e5235 100644 --- a/rust/lancedb/src/dataloader/permutation/split.rs +++ b/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::>(); 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(