This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23740-ed22a4aa19693fccfd47a5beded42d4be1ac3dee in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 1fcdef2fed591e4f5a3144f4d1deb05c57cf343f Author: Fred Thomas <[email protected]> AuthorDate: Tue Jul 21 10:35:24 2026 +0200 `array_agg()` add tests and benchmarks (#23740) ## Which issue does this PR close? - Related to #23715 ## Rationale for this change Add some tests and benchmarks for `array_agg()` to have some measures before performance improvements. ## What changes are included in this PR? 4 unit tests and 2 benchmarks ## Are these changes tested? Changes are only tests and benchmarks, so yes. ## Are there any user-facing changes? No user-facing changes No breaking changes to public APIs --- .../functions-aggregate/benches/array_agg.rs | 105 ++++++++++++++++- datafusion/functions-aggregate/src/array_agg.rs | 124 +++++++++++++++++++++ 2 files changed, 226 insertions(+), 3 deletions(-) diff --git a/datafusion/functions-aggregate/benches/array_agg.rs b/datafusion/functions-aggregate/benches/array_agg.rs index b0d8148c3e..d7e5a51107 100644 --- a/datafusion/functions-aggregate/benches/array_agg.rs +++ b/datafusion/functions-aggregate/benches/array_agg.rs @@ -20,11 +20,14 @@ use std::sync::Arc; use arrow::array::{ Array, ArrayRef, ArrowPrimitiveType, AsArray, ListArray, NullBufferBuilder, + StringArray, }; -use arrow::datatypes::{Field, Int64Type}; +use arrow::datatypes::{DataType, Field, Int64Type}; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_expr::Accumulator; -use datafusion_functions_aggregate::array_agg::ArrayAggAccumulator; +use datafusion_functions_aggregate::array_agg::{ + ArrayAggAccumulator, DistinctArrayAggAccumulator, +}; use arrow::buffer::OffsetBuffer; use arrow::util::bench_util::create_primitive_array; @@ -191,5 +194,101 @@ fn array_agg_benchmark(c: &mut Criterion) { ); } -criterion_group!(benches, array_agg_benchmark); +/// A realistic pool of database names with variable lengths. +const DB_NAMES: &[&str] = &[ + "postgres", + "mysql", + "oracle", + "mssql", + "mongodb", + "redis", + "elasticsearch", + "cassandra", + "dynamodb", + "bigquery", + "snowflake", + "redshift", + "databricks", + "clickhouse", + "duckdb", + "cockroachdb", + "tidb", + "mariadb", + "sqlite", + "neo4j", + "influxdb", + "timescaledb", + "yugabytedb", + "planetscale", + "singlestore", +]; + +/// Low-cardinality: every row is drawn uniformly from `DB_NAMES` (~25 distinct +/// values across 8 192 rows). Exercises the hot duplicate path. +fn create_string_array_low_cardinality(size: usize) -> StringArray { + let mut rng = StdRng::seed_from_u64(42); + StringArray::from_iter_values( + (0..size).map(|_| DB_NAMES[rng.random_range(0..DB_NAMES.len())]), + ) +} + +/// High-cardinality: `db_name_pct` fraction of rows are drawn from `DB_NAMES`; +/// the rest are near-unique random hex strings ("id_XXXXXXXX"). +/// With 8 192 rows and a 32-bit space the collision probability among the +/// random strings is < 1 %, giving ~7 800 distinct values in total. +fn create_string_array_high_cardinality(size: usize, db_name_pct: f32) -> StringArray { + let mut rng = StdRng::seed_from_u64(42); + let strings: Vec<String> = (0..size) + .map(|_| { + if rng.random::<f32>() < db_name_pct { + DB_NAMES[rng.random_range(0..DB_NAMES.len())].to_string() + } else { + format!("id_{:08x}", rng.random::<u32>()) + } + }) + .collect(); + StringArray::from_iter_values(strings.iter().map(String::as_str)) +} + +fn distinct_update_batch_bench( + c: &mut Criterion, + name: &str, + values: &ArrayRef, + ignore_nulls: bool, +) { + c.bench_function(name, |b| { + b.iter(|| { + DistinctArrayAggAccumulator::try_new(&DataType::Utf8, None, ignore_nulls) + .unwrap() + .update_batch(std::slice::from_ref(values)) + .unwrap() + }) + }); +} + +fn distinct_array_agg_benchmark(c: &mut Criterion) { + // --- Low cardinality: ~25 distinct DB names in 8 192 rows --------------- + // Realistic production scenario: most rows are duplicates, the HashSet + // saturates quickly and the rest of the batch is pure dedup overhead. + let values = Arc::new(create_string_array_low_cardinality(8192)) as ArrayRef; + distinct_update_batch_bench( + c, + "distinct_array_agg utf8 low cardinality (~25 distinct)", + &values, + false, + ); + + // --- High cardinality: ~5 % DB names, ~95 % near-unique random strings -- + // Worst-case scenario: almost every row is a new distinct value, so the + // accumulator pays the full insertion cost for nearly every row. + let values = Arc::new(create_string_array_high_cardinality(8192, 0.05)) as ArrayRef; + distinct_update_batch_bench( + c, + "distinct_array_agg utf8 high cardinality (~7800 distinct, 5% db names)", + &values, + false, + ); +} + +criterion_group!(benches, array_agg_benchmark, distinct_array_agg_benchmark); criterion_main!(benches); diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 1dd111f918..1937f17973 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -2589,4 +2589,128 @@ mod tests { Ok(()) } + + #[test] + fn distinct_array_agg_utf8_deduplicates() -> Result<()> { + use arrow::array::StringArray; + + // 7 rows with 4 distinct values, each duplicate appearing twice. + let input: ArrayRef = Arc::new(StringArray::from(vec![ + "postgres", "mysql", "postgres", "redis", "mysql", "duckdb", "redis", + ])); + + let mut acc = DistinctArrayAggAccumulator::try_new(&DataType::Utf8, None, false)?; + acc.update_batch(&[input])?; + + let result = acc.evaluate()?; + let ScalarValue::List(arr) = &result else { + panic!("expected ScalarValue::List, got {result:?}"); + }; + + let inner = arr.value(0); + let strings = inner + .as_any() + .downcast_ref::<StringArray>() + .expect("inner array should be StringArray"); + + // HashSet ordering is nondeterministic — sort before asserting. + let mut values: Vec<&str> = + (0..strings.len()).map(|i| strings.value(i)).collect(); + values.sort_unstable(); + + assert_eq!(values, vec!["duckdb", "mysql", "postgres", "redis"]); + Ok(()) + } + + #[test] + fn distinct_array_agg_int64_deduplicates() -> Result<()> { + use arrow::array::Int64Array; + + // 7 rows with 4 distinct values, each duplicate appearing twice. + let input: ArrayRef = Arc::new(Int64Array::from(vec![1i64, 2, 1, 3, 2, 4, 3])); + + let mut acc = + DistinctArrayAggAccumulator::try_new(&DataType::Int64, None, false)?; + acc.update_batch(&[input])?; + + let result = acc.evaluate()?; + let ScalarValue::List(arr) = &result else { + panic!("expected ScalarValue::List, got {result:?}"); + }; + + let inner = arr.value(0); + let ints = inner + .as_any() + .downcast_ref::<Int64Array>() + .expect("inner array should be Int64Array"); + + let mut values: Vec<i64> = (0..ints.len()).map(|i| ints.value(i)).collect(); + values.sort_unstable(); + + assert_eq!(values, vec![1i64, 2, 3, 4]); + Ok(()) + } + + #[test] + fn distinct_array_agg_float64_deduplicates() -> Result<()> { + use arrow::array::Float64Array; + + // 7 rows with 4 distinct values, each duplicate appearing twice. + let input: ArrayRef = Arc::new(Float64Array::from(vec![ + 1.0f64, 2.5, 1.0, 3.75, 2.5, 4.0, 3.75, + ])); + + let mut acc = + DistinctArrayAggAccumulator::try_new(&DataType::Float64, None, false)?; + acc.update_batch(&[input])?; + + let result = acc.evaluate()?; + let ScalarValue::List(arr) = &result else { + panic!("expected ScalarValue::List, got {result:?}"); + }; + + let inner = arr.value(0); + let floats = inner + .as_any() + .downcast_ref::<Float64Array>() + .expect("inner array should be Float64Array"); + + // f64 has no Ord — use total_cmp for a stable sort. + let mut values: Vec<f64> = (0..floats.len()).map(|i| floats.value(i)).collect(); + values.sort_unstable_by(|a, b| a.total_cmp(b)); + + assert_eq!(values, vec![1.0f64, 2.5, 3.75, 4.0]); + Ok(()) + } + + #[test] + fn distinct_array_agg_date32_deduplicates() -> Result<()> { + use arrow::array::Date32Array; + + // 7 rows with 4 distinct dates (days since epoch), each duplicate appearing twice. + let input: ArrayRef = Arc::new(Date32Array::from(vec![ + 100i32, 200, 100, 300, 200, 400, 300, + ])); + + let mut acc = + DistinctArrayAggAccumulator::try_new(&DataType::Date32, None, false)?; + acc.update_batch(&[input])?; + + let result = acc.evaluate()?; + let ScalarValue::List(arr) = &result else { + panic!("expected ScalarValue::List, got {result:?}"); + }; + + let inner = arr.value(0); + let dates = inner + .as_any() + .downcast_ref::<Date32Array>() + .expect("inner array should be Date32Array"); + + let mut values: Vec<i32> = (0..dates.len()).map(|i| dates.value(i)).collect(); + values.sort_unstable(); + + assert_eq!(values, vec![100i32, 200, 300, 400]); + Ok(()) + } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
