diff --git a/datafusion/fuzz/tests/fuzz.rs b/datafusion/fuzz/tests/fuzz.rs index 6eae5f9127eb5..fd4a0ff1a00d7 100644 --- a/datafusion/fuzz/tests/fuzz.rs +++ b/datafusion/fuzz/tests/fuzz.rs @@ -20,15 +20,27 @@ mod suites; #[tokio::test(flavor = "multi_thread")] async fn fuzz_tests() -> Result<()> { let spill_pool = suites::spill_pool::SpillPoolFuzzer; - let suites = [Suite::new( - &spill_pool, - RunConfig { - run_seed: Some(0), - max_cases: Some(1000), - time_limit: None, - cases_per_dataset: 1000, - }, - )]; + let key_range_bitmap = suites::key_range_bitmap::KeyRangeBitmapFuzzer; + let suites = [ + Suite::new( + &spill_pool, + RunConfig { + run_seed: Some(0), + max_cases: Some(1000), + time_limit: None, + cases_per_dataset: 1000, + }, + ), + Suite::new( + &key_range_bitmap, + RunConfig { + run_seed: Some(0), + max_cases: Some(10000), + time_limit: None, + cases_per_dataset: 10000, + }, + ), + ]; run_suites(suites).await } diff --git a/datafusion/fuzz/tests/suites/key_range_bitmap.rs b/datafusion/fuzz/tests/suites/key_range_bitmap.rs new file mode 100644 index 0000000000000..c0dbebcd093a2 --- /dev/null +++ b/datafusion/fuzz/tests/suites/key_range_bitmap.rs @@ -0,0 +1,195 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Randomized partitions against [`KeyRangeBitmap::union`]'s contract: it +//! never excludes a key any part may hold, equals the bitmap built from all +//! parts' keys when no part brings its own bitmap, and only widens that +//! bitmap's buckets when one does. + +use std::future::Future; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array}; +use arrow::buffer::BooleanBuffer; +use datafusion_common::{Result, ScalarValue}; +use datafusion_physical_plan::joins::key_range_bitmap::{KeyRangeBitmap, KeyRangePart}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use test_utils::fuzz::FuzzTest; + +pub(crate) struct KeyRangeBitmapFuzzer; + +fn parts_of(b: &KeyRangeBitmap) -> (u64, u64, u32, &BooleanBuffer) { + (b.offset(), b.span(), b.shift(), b.bits()) +} + +enum Part { + Keys(ArrayRef), + Range(ScalarValue, ScalarValue), + Bitmap(KeyRangeBitmap), +} + +impl FuzzTest for KeyRangeBitmapFuzzer { + type Dataset = (); + + fn name(&self) -> &str { + "key_range_bitmap_union" + } + + fn prepare_dataset(&self, _seed: u64) -> impl Future> { + std::future::ready(Ok(())) + } + + fn run_case(&self, _dataset: &(), seed: u64) -> impl Future> { + let mut rng = StdRng::seed_from_u64(seed); + // Domain widths from one key per bucket to many keys per bucket, so + // the union's buckets are sometimes exact and sometimes coarse. + let width = [64i64, 1 << 12, 1 << 24, 1 << 40][rng.random_range(0..4)]; + let start = rng.random_range(-(1i64 << 50)..(1i64 << 50)); + let key = |rng: &mut StdRng| start + rng.random_range(0..width); + let keys_array = + |keys: &[i64]| -> ArrayRef { Arc::new(Int64Array::from(keys.to_vec())) }; + let scalar = |key: i64| ScalarValue::Int64(Some(key)); + + // Each part, and every key any part may hold. + let mut parts = vec![]; + let mut all_keys = vec![]; + let mut has_bitmap = false; + for _ in 0..rng.random_range(1..5) { + if rng.random_bool(1.0 / 3.0) { + // Narrow, so every key it may hold can be listed. + let lo = key(&mut rng); + let hi = lo + rng.random_range(0..64); + parts.push(Part::Range(scalar(lo), scalar(hi))); + all_keys.extend(lo..=hi); + continue; + } + let keys: Vec = (0..rng.random_range(1..50)) + .map(|_| key(&mut rng)) + .collect(); + let (lo, hi) = (*keys.iter().min().unwrap(), *keys.iter().max().unwrap()); + let bitmap = rng + .random_bool(0.5) + .then(|| { + KeyRangeBitmap::try_new( + &keys_array(&keys), + &scalar(lo), + &scalar(hi), + 0, + ) + }) + .flatten(); + parts.push(match bitmap { + Some(bitmap) => { + has_bitmap = true; + Part::Bitmap(bitmap) + } + None => Part::Keys(keys_array(&keys)), + }); + all_keys.extend(keys); + } + let min = ScalarValue::Int64(all_keys.iter().min().copied()); + let max = ScalarValue::Int64(all_keys.iter().max().copied()); + let context = format!( + "seed {seed}: width {width}, {} parts, {} keys", + parts.len(), + all_keys.len() + ); + + let union = KeyRangeBitmap::union( + &min, + &max, + &parts + .iter() + .map(|part| match part { + Part::Keys(keys) => KeyRangePart::Keys(keys.as_ref()), + Part::Range(min, max) => KeyRangePart::Range { min, max }, + Part::Bitmap(bitmap) => KeyRangePart::Bitmap(bitmap), + }) + .collect::>(), + ); + // Reference: one bitmap built straight from every key any part may hold. + let exact = KeyRangeBitmap::try_new(&keys_array(&all_keys), &min, &max, 0); + + // Never excludes a key any part may hold. + if let Some(union) = &union { + for key in &all_keys { + assert!( + union.may_contain_range(*key as u64, *key as u64), + "{context}: dropped key {key}" + ); + } + } + if has_bitmap { + // A part's bitmap can only widen buckets: every bucket the exact + // bitmap sets, the union sets too. + match (&union, &exact) { + (_, None) => assert!(union.is_none(), "{context}"), + (None, Some(_)) => {} + (Some(union), Some(exact)) => { + assert_eq!( + (union.offset(), union.span(), union.shift()), + (exact.offset(), exact.span(), exact.shift()), + "{context}" + ); + for bucket in exact.bits().set_indices() { + assert!(union.bits().value(bucket), "{context}: bucket {bucket}"); + } + } + } + // Each set bucket of a part's bitmap stands for its whole key range, + // so the union equals the one over those ranges given as `Range`s. + let ranges: Vec<(ScalarValue, ScalarValue)> = parts + .iter() + .flat_map(|part| match part { + Part::Bitmap(bitmap) => bitmap + .bits() + .set_indices() + .map(|bucket| { + let lo = (bucket as u64) << bitmap.shift(); + let hi = lo + .saturating_add((1u64 << bitmap.shift()) - 1) + .min(bitmap.span()); + let key = + |at: u64| scalar(bitmap.offset().wrapping_add(at) as i64); + (key(lo), key(hi)) + }) + .collect(), + _ => vec![], + }) + .collect(); + let expanded = KeyRangeBitmap::union( + &min, + &max, + &parts + .iter() + .filter_map(|part| match part { + Part::Keys(keys) => Some(KeyRangePart::Keys(keys.as_ref())), + Part::Range(min, max) => Some(KeyRangePart::Range { min, max }), + Part::Bitmap(_) => None, + }) + .chain( + ranges + .iter() + .map(|(min, max)| KeyRangePart::Range { min, max }), + ) + .collect::>(), + ); + assert_eq!( + union.as_ref().map(parts_of), + expanded.as_ref().map(parts_of), + "{context}: bitmap parts differ from their bucket ranges" + ); + } else { + // Keys and ranges mark exactly the buckets their keys fall in. + assert_eq!( + union.as_ref().map(parts_of), + exact.as_ref().map(parts_of), + "{context}" + ); + } + + std::future::ready(Ok(())) + } +} diff --git a/datafusion/fuzz/tests/suites/mod.rs b/datafusion/fuzz/tests/suites/mod.rs index a443d433bd089..5edf91e5d658b 100644 --- a/datafusion/fuzz/tests/suites/mod.rs +++ b/datafusion/fuzz/tests/suites/mod.rs @@ -4,4 +4,5 @@ //! Implementation for fuzz test cases. +pub(crate) mod key_range_bitmap; pub(crate) mod spill_pool; diff --git a/datafusion/physical-plan/src/joins/key_range_bitmap.rs b/datafusion/physical-plan/src/joins/key_range_bitmap.rs index 47030f5cadc4a..2ad9c3369e53a 100644 --- a/datafusion/physical-plan/src/joins/key_range_bitmap.rs +++ b/datafusion/physical-plan/src/joins/key_range_bitmap.rs @@ -16,7 +16,12 @@ use crate::joins::array_map::ArrayMap; const MAX_BUCKETS: u64 = 1 << 20; /// One build partition's contribution to [`KeyRangeBitmap::union`]. -pub(crate) enum KeyRangePart<'a> { +/// +/// # Public Only for Internal Use: +/// The `datafusion-fuzz` tests require this, but it is not part of the +/// supported public API. +#[doc(hidden)] +pub enum KeyRangePart<'a> { /// The partition's own bitmap. Bitmap(&'a KeyRangeBitmap), /// The partition's build keys, which must lie within its `[min, max]`. @@ -107,7 +112,12 @@ impl KeyRangeBitmap { /// Coarser than the parts' own bitmaps where its buckets are wider, but /// never clears a bucket a part may hold a key in. `None` when the key /// type has no `u64` ordering, or when nothing can be excluded. - pub(crate) fn union( + /// + /// # Public Only for Internal Use: + /// The `datafusion-fuzz` tests require this, but it is not part of the + /// supported public API. + #[doc(hidden)] + pub fn union( min: &ScalarValue, max: &ScalarValue, parts: &[KeyRangePart<'_>], @@ -252,6 +262,32 @@ impl KeyRangeBitmap { } } +/// # Public Only for Internal Use: +/// The `datafusion-fuzz` tests require these accessors, but they are not part +/// of the supported public API. +#[doc(hidden)] +impl KeyRangeBitmap { + /// Smallest key, in the `u64` ordering the join's array map uses. + pub fn offset(&self) -> u64 { + self.offset + } + + /// `max - min`: a key is in range iff `key.wrapping_sub(offset) <= span`. + pub fn span(&self) -> u64 { + self.span + } + + /// Each bucket covers `1 << shift` keys. + pub fn shift(&self) -> u32 { + self.shift + } + + /// One bit per bucket. + pub fn bits(&self) -> &BooleanBuffer { + &self.bits + } +} + /// One bit per bucket, set for each bucket holding a non-null key of `keys`. fn bucket_bits(keys: &dyn Array, offset: u64, shift: u32, buckets: u64) -> BooleanBuffer { let mut words = MutableBuffer::new_null(buckets as usize); @@ -425,189 +461,4 @@ mod tests { assert!(may(&s, -1, 1)); assert!(!may(&s, 1 << 50, 1 << 51)); } - - /// `(offset, span, shift, bits)`, to compare bitmaps by value. - fn parts_of(b: &KeyRangeBitmap) -> (u64, u64, u32, &BooleanBuffer) { - (b.offset, b.span, b.shift, &b.bits) - } - - /// One randomly generated scenario, see [`union_fuzz`]. - fn run_union_scenario(seed: u64) { - use rand::rngs::StdRng; - use rand::{Rng, SeedableRng}; - - let mut rng = StdRng::seed_from_u64(seed); - // Domain widths from one key per bucket to many keys per bucket, so - // the union's buckets are sometimes exact and sometimes coarse. - let width = [64i64, 1 << 12, 1 << 24, 1 << 40][rng.random_range(0..4)]; - let start = rng.random_range(-(1i64 << 50)..(1i64 << 50)); - let key = |rng: &mut StdRng| start + rng.random_range(0..width); - let keys_array = - |keys: &[i64]| -> ArrayRef { Arc::new(Int64Array::from(keys.to_vec())) }; - let scalar = |key: i64| ScalarValue::Int64(Some(key)); - - enum Part { - Keys(ArrayRef), - Range(ScalarValue, ScalarValue), - Bitmap(KeyRangeBitmap), - } - // Each part, and every key any part may hold. - let mut parts = vec![]; - let mut all_keys = vec![]; - let mut has_bitmap = false; - for _ in 0..rng.random_range(1..5) { - if rng.random_bool(1.0 / 3.0) { - // Narrow, so every key it may hold can be listed. - let lo = key(&mut rng); - let hi = lo + rng.random_range(0..64); - parts.push(Part::Range(scalar(lo), scalar(hi))); - all_keys.extend(lo..=hi); - continue; - } - let keys: Vec = (0..rng.random_range(1..50)) - .map(|_| key(&mut rng)) - .collect(); - let (lo, hi) = (*keys.iter().min().unwrap(), *keys.iter().max().unwrap()); - let bitmap = rng - .random_bool(0.5) - .then(|| { - KeyRangeBitmap::try_new( - &keys_array(&keys), - &scalar(lo), - &scalar(hi), - 0, - ) - }) - .flatten(); - parts.push(match bitmap { - Some(bitmap) => { - has_bitmap = true; - Part::Bitmap(bitmap) - } - None => Part::Keys(keys_array(&keys)), - }); - all_keys.extend(keys); - } - let min = ScalarValue::Int64(all_keys.iter().min().copied()); - let max = ScalarValue::Int64(all_keys.iter().max().copied()); - let context = format!( - "seed {seed}: width {width}, {} parts, {} keys", - parts.len(), - all_keys.len() - ); - - let union = KeyRangeBitmap::union( - &min, - &max, - &parts - .iter() - .map(|part| match part { - Part::Keys(keys) => KeyRangePart::Keys(keys.as_ref()), - Part::Range(min, max) => KeyRangePart::Range { min, max }, - Part::Bitmap(bitmap) => KeyRangePart::Bitmap(bitmap), - }) - .collect::>(), - ); - // Reference: one bitmap built straight from every key any part may hold. - let exact = KeyRangeBitmap::try_new(&keys_array(&all_keys), &min, &max, 0); - - // Never excludes a key any part may hold. - if let Some(union) = &union { - for key in &all_keys { - assert!(may(union, *key, *key), "{context}: dropped key {key}"); - } - } - if has_bitmap { - // A part's bitmap can only widen buckets: every bucket the exact - // bitmap sets, the union sets too. - match (&union, &exact) { - (_, None) => assert!(union.is_none(), "{context}"), - (None, Some(_)) => {} - (Some(union), Some(exact)) => { - assert_eq!( - (union.offset, union.span, union.shift), - (exact.offset, exact.span, exact.shift), - "{context}" - ); - for bucket in exact.bits.set_indices() { - assert!(union.bits.value(bucket), "{context}: bucket {bucket}"); - } - } - } - // Each set bucket of a part's bitmap stands for its whole key range, - // so the union equals the one over those ranges given as `Range`s. - let ranges: Vec<(ScalarValue, ScalarValue)> = parts - .iter() - .flat_map(|part| match part { - Part::Bitmap(bitmap) => bitmap - .bits - .set_indices() - .map(|bucket| { - let lo = (bucket as u64) << bitmap.shift; - let hi = lo - .saturating_add((1u64 << bitmap.shift) - 1) - .min(bitmap.span); - let key = - |at: u64| scalar(bitmap.offset.wrapping_add(at) as i64); - (key(lo), key(hi)) - }) - .collect(), - _ => vec![], - }) - .collect(); - let expanded = KeyRangeBitmap::union( - &min, - &max, - &parts - .iter() - .filter_map(|part| match part { - Part::Keys(keys) => Some(KeyRangePart::Keys(keys.as_ref())), - Part::Range(min, max) => Some(KeyRangePart::Range { min, max }), - Part::Bitmap(_) => None, - }) - .chain( - ranges - .iter() - .map(|(min, max)| KeyRangePart::Range { min, max }), - ) - .collect::>(), - ); - assert_eq!( - union.as_ref().map(parts_of), - expanded.as_ref().map(parts_of), - "{context}: bitmap parts differ from their bucket ranges" - ); - } else { - // Keys and ranges mark exactly the buckets their keys fall in. - assert_eq!( - union.as_ref().map(parts_of), - exact.as_ref().map(parts_of), - "{context}" - ); - } - } - - /// Randomized partitions against [`KeyRangeBitmap::union`]'s contract: it - /// never excludes a key any part may hold, equals the bitmap built from all - /// parts' keys when no part brings its own bitmap, and only widens that - /// bitmap's buckets when one does. - /// - /// `DATAFUSION_KEY_RANGE_UNION_FUZZ_ITERATIONS` and - /// `DATAFUSION_KEY_RANGE_UNION_FUZZ_SEED` select how many scenarios run and - /// from which seed; a failure names its seed so it can be replayed. - #[test] - fn union_fuzz() { - let env_u64 = |name: &str, default: u64| { - std::env::var(name) - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(default) - }; - let first_seed = env_u64("DATAFUSION_KEY_RANGE_UNION_FUZZ_SEED", 0); - let iterations = env_u64("DATAFUSION_KEY_RANGE_UNION_FUZZ_ITERATIONS", 256); - // `saturating_add` so a replay seed near `u64::MAX` cannot overflow. - for seed in first_seed..first_seed.saturating_add(iterations) { - run_union_scenario(seed); - } - } } diff --git a/xtask/src/ci_steps.rs b/xtask/src/ci_steps.rs index a4cfc8f831fc8..e663688310343 100644 --- a/xtask/src/ci_steps.rs +++ b/xtask/src/ci_steps.rs @@ -16,10 +16,8 @@ const CI_COMMAND: &str = "cargo xtask ci step"; const CI_SHORTCUT: &str = "cargo ci-step"; const TPCH_DATA_DIR: &str = "datafusion/sqllogictest/test_files/tpch/data"; /// Physical-plan fuzzers awaiting migration to the centralized driver. -const LEGACY_FUZZ_TESTS: &[&str] = &[ - "joins::hash_join::partitioned_hash_eval::tests::routed_lookup_fuzz", - "joins::key_range_bitmap::tests::union_fuzz", -]; +const LEGACY_FUZZ_TESTS: &[&str] = + &["joins::hash_join::partitioned_hash_eval::tests::routed_lookup_fuzz"]; /// Metadata and implementation for one CI step. /// @@ -498,8 +496,7 @@ impl StepContext { ]) .args(LEGACY_FUZZ_TESTS) .env("RUST_BACKTRACE", "1") - .env("DATAFUSION_ROUTED_LOOKUP_FUZZ_ITERATIONS", "10000") - .env("DATAFUSION_KEY_RANGE_UNION_FUZZ_ITERATIONS", "10000"); + .env("DATAFUSION_ROUTED_LOOKUP_FUZZ_ITERATIONS", "10000"); } "hash-collisions" => { cmd_builder.doc( @@ -990,14 +987,13 @@ cargo test \ # Run unmigrated physical-plan fuzz tests with more iterations. cd /workspace && \ -DATAFUSION_KEY_RANGE_UNION_FUZZ_ITERATIONS=10000 \ DATAFUSION_ROUTED_LOOKUP_FUZZ_ITERATIONS=10000 \ RUST_BACKTRACE=1 \ cargo test \ --profile ci \ -p datafusion-physical-plan \ --lib \ --- --exact joins::hash_join::partitioned_hash_eval::tests::routed_lookup_fuzz joins::key_range_bitmap::tests::union_fuzz +-- --exact joins::hash_join::partitioned_hash_eval::tests::routed_lookup_fuzz ") }, },