mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-14 14:46:53 +00:00
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) <noreply@anthropic.com>
This commit is contained in:
Will Jones
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
7ebfc71247
commit
6858a9370f
@@ -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 = [
|
||||
|
||||
@@ -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()
|
||||
)
|
||||
|
||||
|
||||
Generated
+23
-8
@@ -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" },
|
||||
|
||||
@@ -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::<Vec<_>>()
|
||||
.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)
|
||||
}
|
||||
|
||||
@@ -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::<UInt64Type>().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%
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user