From 6858a9370fcfe993787eca50c309b5de9b8fba03 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Wed, 15 Jul 2026 15:09:19 -0700 Subject: [PATCH] fix: adapt permutation tests and datafusion python to lance v9 DataFusion 54 (pulled in by lance v9) replaced its hashing RandomState (ahash -> foldhash), which changes hash-based split output. Update the hardcoded expectations in test_hash_split, and hash a high-cardinality column in the Python test_split_hash_with_discard so its discard behavior is stable. The new lance file reader returns no batches for an empty file (the old read_record_batch returned a single empty batch), so concat_batches on batches[0] panicked when a shuffle output file was empty. Build the batch from the known schema instead, which handles the empty case. Bump the datafusion Python package to v54 to match the bundled DataFusion, avoiding an FFI ABI mismatch in register_table_provider. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/pyproject.toml | 2 +- python/python/tests/test_permutation.py | 5 ++- python/uv.lock | 31 ++++++++++++++----- .../src/dataloader/permutation/shuffle.rs | 6 +++- .../src/dataloader/permutation/split.rs | 8 ++--- 5 files changed, 37 insertions(+), 15 deletions(-) diff --git a/python/pyproject.toml b/python/pyproject.toml index 4e8391787..bfbb30275 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -65,7 +65,7 @@ tests = [ "pyarrow-stubs>=16.0", "pylance>=7,<8", "requests>=2.31.0", - "datafusion>=53,<54", + "datafusion>=54,<55", "opentelemetry-sdk>=1.30.0", ] dev = [ diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index 0da12c7dc..7ff35e5e2 100644 --- a/python/python/tests/test_permutation.py +++ b/python/python/tests/test_permutation.py @@ -134,8 +134,11 @@ def test_split_hash_with_discard(mem_db): ) permutation_tbl = ( + # Hash a high-cardinality column: "category" has only two distinct + # values, so whether anything is discarded would hinge on where those + # two hashes land rather than on the discard weight. permutation_builder(tbl) - .split_hash(["category"], [1, 1], discard_weight=2) # Should discard ~50% + .split_hash(["id"], [1, 1], discard_weight=2) # Should discard ~50% .execute() ) diff --git a/python/uv.lock b/python/uv.lock index 7e6c9b5e6..973e7287c 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -657,6 +657,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/0d/67e5b4109ea4a837e80daa87c2c696711955e40449a97e8926672534def2/click-8.4.1-py3-none-any.whl", hash = "sha256:482be17c6991b8c19c5429a1e995d9b0efdbb63172824c41f99965dc0ade8ec2", size = 116639, upload-time = "2026-05-22T04:08:35.26Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "cohere" version = "7.0.3" @@ -850,19 +859,25 @@ nvtx = [ [[package]] name = "datafusion" -version = "53.0.0" +version = "54.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "cloudpickle" }, { name = "pyarrow" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/58/2b/0f96f12b70839c93930c4e17d767fc32b6c77d548c78784128049e944701/datafusion-53.0.0.tar.gz", hash = "sha256:ba9a5ec06b5453fbd8710d6aeeb515a8bcac4b6c140e254409bb53a5f322ef22", size = 224267, upload-time = "2026-04-13T00:45:02.686Z" } +sdist = { url = "https://files.pythonhosted.org/packages/60/90/886f7e9cf827f07ebd60bd293e54e0a028a50dd49bbaef0ee42aae1981ea/datafusion-54.0.0.tar.gz", hash = "sha256:cfe7e8dfc026efc05824f49b53ad6a72caf5c2d6820759b6212a09e245a427ed", size = 276448, upload-time = "2026-06-29T11:19:34.816Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/af/4c/60e052813d81f1ffe3123ead013dbdd2cf961daa576cb9056cbb80228e6b/datafusion-53.0.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:a0bd1a98d736571321416dc4ed361a9d1225da1ec9f6c5fad818d75f547697a7", size = 35774913, upload-time = "2026-04-13T00:44:46.235Z" }, - { url = "https://files.pythonhosted.org/packages/6e/59/beabe5301df3338d8206446cd624079e43bdad46e20377a6336017fb6ccf/datafusion-53.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:ce186a8d2405afd67e11e2fb75715019f16b00d070b8d0da89d8aa61cc74c8b5", size = 32667118, upload-time = "2026-04-13T00:44:50.269Z" }, - { url = "https://files.pythonhosted.org/packages/ae/94/636ab61ade98395daea6e733e225e9c7beef111c7c5b575ac851513e203c/datafusion-53.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:288a00a7ef03e2807a4667683f7560efd80d60ed1d41696ac15ca9ded14c8251", size = 35585824, upload-time = "2026-04-13T00:44:53.683Z" }, - { url = "https://files.pythonhosted.org/packages/34/80/b9f4889209af02f8d14bccb0e6f0519c329b072bc4d2595025a1303f144c/datafusion-53.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:8fef0004f0161fcfc556c025a7201f9cc3169aa3adb97a86419ebb34182d9efb", size = 38083690, upload-time = "2026-04-13T00:44:57.188Z" }, - { url = "https://files.pythonhosted.org/packages/4b/1a/ea4831fc6aeefedbcf186c9f6a273d507b1787c03cbb905bded7e1149a6a/datafusion-53.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:4c8410f5f659b926677be6c7d443bbc05d825c078c970b7d8cf977ebcf948314", size = 38120687, upload-time = "2026-04-13T00:45:00.633Z" }, + { url = "https://files.pythonhosted.org/packages/46/58/4c5b981e3d9ade32a906c15a4941eef50c9b862781cdc14bf4dff48d026a/datafusion-54.0.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:946f55e48b8d523d7b4ac106bdf588b4493c2c66f81877d6952aafeaf7c3ec73", size = 39810553, upload-time = "2026-06-29T11:19:02.1Z" }, + { url = "https://files.pythonhosted.org/packages/66/e5/5e4dbd42ce9a2affb3be90d9ab17cebde1a6f28b0d9fb4b83d612d5c8e42/datafusion-54.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:2a3bf43185c7e43e25242e5fb17b6a11b86bf976434c0bc493fdedbd9a080969", size = 37145255, upload-time = "2026-06-29T11:19:05.491Z" }, + { url = "https://files.pythonhosted.org/packages/c6/5e/dbb9e6e3e5006d34f295d7ac73f1302c8f2df140666402a06e6c55028edb/datafusion-54.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:9432bf162381e9282cbc74915b8b773895de18be836f7e3f6d0de4d981f24630", size = 38853856, upload-time = "2026-06-29T11:19:08.732Z" }, + { url = "https://files.pythonhosted.org/packages/a8/81/e69008e3479f4d0134875bc4ae39503bedcd55ca2597e71392c963c651b4/datafusion-54.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3bcd4d213fa74710e75e6e182cc468c2bdbc5ffc74a08c8155d414fbbfa1b3f6", size = 41050149, upload-time = "2026-06-29T11:19:12.108Z" }, + { url = "https://files.pythonhosted.org/packages/61/d4/8ba6e3fe3291c9ccc94b5ca3ec3c1fbcbfbe5ece5ffb965e4550844e2c56/datafusion-54.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:b934e097e1bdca7d5768a81ac1bc4a1812cb459269f8b1a5d892a5d930f18376", size = 43444869, upload-time = "2026-06-29T11:19:15.963Z" }, + { url = "https://files.pythonhosted.org/packages/9d/41/5608323226f21a0fa180823c531dbc0ed270e9b694f299b7647505cb6a06/datafusion-54.0.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:c4e79048da82ad89b768bd0be7df39254cd2a0afe2b719d1f129e8a7229af683", size = 39796248, upload-time = "2026-06-29T11:19:19.208Z" }, + { url = "https://files.pythonhosted.org/packages/18/81/392ee323104ab14ca689384723b69e137064a828233c165574f97a74c0e9/datafusion-54.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:fe57038003b18e28b90752c1e32b44af74ec4f552a1904aee725e1129a00c447", size = 37153577, upload-time = "2026-06-29T11:19:22.397Z" }, + { url = "https://files.pythonhosted.org/packages/40/c4/ebd5ef5349ecbea7f5f9da76c213581c13e7bbe1b5735c9925b279eeb4eb/datafusion-54.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:574f642832a106456cfc4f32aa82484c504fc32f4be2b510202bcb579de8e6d1", size = 38849839, upload-time = "2026-06-29T11:19:25.783Z" }, + { url = "https://files.pythonhosted.org/packages/5a/b9/2383d30d317bb913cab97dbf2e6e1d5f37f594860d5c5bc176e025cf7d4a/datafusion-54.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:796fd5683927443c5bc61999d00b9007ef9b5ce107725ea8d241df718860985d", size = 41074623, upload-time = "2026-06-29T11:19:29.119Z" }, + { url = "https://files.pythonhosted.org/packages/35/5c/553fd1107dede0a56727fda7216a7198d41394f2d19697f4fb104cc695ea/datafusion-54.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:64973c63874ec31670dd97b32b18af7b07fad679cb20d58ed154038e3a5c204e", size = 43438801, upload-time = "2026-06-29T11:19:32.799Z" }, ] [[package]] @@ -1951,7 +1966,7 @@ requires-dist = [ { name = "botocore", marker = "extra == 'embeddings'", specifier = ">=1.31.57" }, { name = "cohere", marker = "extra == 'embeddings'", specifier = ">=4.0" }, { name = "colpali-engine", marker = "extra == 'embeddings'", specifier = ">=0.3.10" }, - { name = "datafusion", marker = "extra == 'tests'", specifier = ">=53,<54" }, + { name = "datafusion", marker = "extra == 'tests'", specifier = ">=54,<55" }, { name = "deprecation", specifier = ">=2.1.0" }, { name = "duckdb", marker = "extra == 'tests'", specifier = ">=0.9.0" }, { name = "google-genai", marker = "extra == 'embeddings'", specifier = ">=1.0.0" }, diff --git a/rust/lancedb/src/dataloader/permutation/shuffle.rs b/rust/lancedb/src/dataloader/permutation/shuffle.rs index 370149331..cf118a996 100644 --- a/rust/lancedb/src/dataloader/permutation/shuffle.rs +++ b/rust/lancedb/src/dataloader/permutation/shuffle.rs @@ -216,6 +216,7 @@ impl Shuffler { let scan_scheduler = ScanScheduler::new(Arc::new(object_store), scheduler_config); let job_id = self.id.clone(); let rng = Arc::new(Mutex::new(rng)); + let read_schema = arrow_schema.clone(); // Second pass, read each file as a single batch and shuffle let stream = futures::stream::iter(0..num_files) @@ -224,6 +225,7 @@ impl Shuffler { let rng = rng.clone(); let tmp_dir = tmp_dir.clone(); let job_id = job_id.clone(); + let read_schema = read_schema.clone(); async move { let path = tmp_dir.join(format!("shuffle_{}_{file_index}.lance", job_id)); let path = object_store::path::Path::from_absolute_path(path).unwrap(); @@ -249,7 +251,9 @@ impl Shuffler { .await? .try_collect::>() .await?; - let batch = concat_batches(&batches[0].schema(), &batches)?; + // An empty file yields no batches; fall back to an empty batch + // with the expected schema so shuffling handles it gracefully. + let batch = concat_batches(&read_schema, &batches)?; let mut rng = rng.lock().unwrap_or_else(|e| e.into_inner()); Self::shuffle_batch(&batch, &mut rng, clump_size) } diff --git a/rust/lancedb/src/dataloader/permutation/split.rs b/rust/lancedb/src/dataloader/permutation/split.rs index 202691bf1..b33544100 100644 --- a/rust/lancedb/src/dataloader/permutation/split.rs +++ b/rust/lancedb/src/dataloader/permutation/split.rs @@ -796,15 +796,15 @@ mod tests { // These assertions are all based on fixed seed in data generation but they match // up roughly to what we expect (25% discarded, 25% in split 0, 50% in split 1) - // 14 rows (28%) are discarded because discard_weight is 1 - assert_eq!(split_batch.num_rows(), 36); + // 8 rows (16%) are discarded because discard_weight is 1 + assert_eq!(split_batch.num_rows(), 42); assert_eq!(split_batch.num_columns(), 2); let split_ids = split_batch.column(1).as_primitive::().values(); let num_in_split_0 = split_ids.iter().filter(|v| **v == 0).count(); let num_in_split_1 = split_ids.iter().filter(|v| **v == 1).count(); - assert_eq!(num_in_split_0, 11); // 22% - assert_eq!(num_in_split_1, 25); // 50% + assert_eq!(num_in_split_0, 12); // 24% + assert_eq!(num_in_split_1, 30); // 60% } }