This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 1fcdef2fed `array_agg()` add tests and benchmarks (#23740)
1fcdef2fed is described below
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]
