Skip to content

Commit 1fcdef2

Browse files
authored
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
1 parent ed22a4a commit 1fcdef2

2 files changed

Lines changed: 226 additions & 3 deletions

File tree

‎datafusion/functions-aggregate/benches/array_agg.rs‎

Lines changed: 102 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,14 @@ use std::sync::Arc;
2020

2121
use arrow::array::{
2222
Array, ArrayRef, ArrowPrimitiveType, AsArray, ListArray, NullBufferBuilder,
23+
StringArray,
2324
};
24-
use arrow::datatypes::{Field, Int64Type};
25+
use arrow::datatypes::{DataType, Field, Int64Type};
2526
use criterion::{Criterion, criterion_group, criterion_main};
2627
use datafusion_expr::Accumulator;
27-
use datafusion_functions_aggregate::array_agg::ArrayAggAccumulator;
28+
use datafusion_functions_aggregate::array_agg::{
29+
ArrayAggAccumulator, DistinctArrayAggAccumulator,
30+
};
2831

2932
use arrow::buffer::OffsetBuffer;
3033
use arrow::util::bench_util::create_primitive_array;
@@ -191,5 +194,101 @@ fn array_agg_benchmark(c: &mut Criterion) {
191194
);
192195
}
193196

194-
criterion_group!(benches, array_agg_benchmark);
197+
/// A realistic pool of database names with variable lengths.
198+
const DB_NAMES: &[&str] = &[
199+
"postgres",
200+
"mysql",
201+
"oracle",
202+
"mssql",
203+
"mongodb",
204+
"redis",
205+
"elasticsearch",
206+
"cassandra",
207+
"dynamodb",
208+
"bigquery",
209+
"snowflake",
210+
"redshift",
211+
"databricks",
212+
"clickhouse",
213+
"duckdb",
214+
"cockroachdb",
215+
"tidb",
216+
"mariadb",
217+
"sqlite",
218+
"neo4j",
219+
"influxdb",
220+
"timescaledb",
221+
"yugabytedb",
222+
"planetscale",
223+
"singlestore",
224+
];
225+
226+
/// Low-cardinality: every row is drawn uniformly from `DB_NAMES` (~25 distinct
227+
/// values across 8 192 rows). Exercises the hot duplicate path.
228+
fn create_string_array_low_cardinality(size: usize) -> StringArray {
229+
let mut rng = StdRng::seed_from_u64(42);
230+
StringArray::from_iter_values(
231+
(0..size).map(|_| DB_NAMES[rng.random_range(0..DB_NAMES.len())]),
232+
)
233+
}
234+
235+
/// High-cardinality: `db_name_pct` fraction of rows are drawn from `DB_NAMES`;
236+
/// the rest are near-unique random hex strings ("id_XXXXXXXX").
237+
/// With 8 192 rows and a 32-bit space the collision probability among the
238+
/// random strings is < 1 %, giving ~7 800 distinct values in total.
239+
fn create_string_array_high_cardinality(size: usize, db_name_pct: f32) -> StringArray {
240+
let mut rng = StdRng::seed_from_u64(42);
241+
let strings: Vec<String> = (0..size)
242+
.map(|_| {
243+
if rng.random::<f32>() < db_name_pct {
244+
DB_NAMES[rng.random_range(0..DB_NAMES.len())].to_string()
245+
} else {
246+
format!("id_{:08x}", rng.random::<u32>())
247+
}
248+
})
249+
.collect();
250+
StringArray::from_iter_values(strings.iter().map(String::as_str))
251+
}
252+
253+
fn distinct_update_batch_bench(
254+
c: &mut Criterion,
255+
name: &str,
256+
values: &ArrayRef,
257+
ignore_nulls: bool,
258+
) {
259+
c.bench_function(name, |b| {
260+
b.iter(|| {
261+
DistinctArrayAggAccumulator::try_new(&DataType::Utf8, None, ignore_nulls)
262+
.unwrap()
263+
.update_batch(std::slice::from_ref(values))
264+
.unwrap()
265+
})
266+
});
267+
}
268+
269+
fn distinct_array_agg_benchmark(c: &mut Criterion) {
270+
// --- Low cardinality: ~25 distinct DB names in 8 192 rows ---------------
271+
// Realistic production scenario: most rows are duplicates, the HashSet
272+
// saturates quickly and the rest of the batch is pure dedup overhead.
273+
let values = Arc::new(create_string_array_low_cardinality(8192)) as ArrayRef;
274+
distinct_update_batch_bench(
275+
c,
276+
"distinct_array_agg utf8 low cardinality (~25 distinct)",
277+
&values,
278+
false,
279+
);
280+
281+
// --- High cardinality: ~5 % DB names, ~95 % near-unique random strings --
282+
// Worst-case scenario: almost every row is a new distinct value, so the
283+
// accumulator pays the full insertion cost for nearly every row.
284+
let values = Arc::new(create_string_array_high_cardinality(8192, 0.05)) as ArrayRef;
285+
distinct_update_batch_bench(
286+
c,
287+
"distinct_array_agg utf8 high cardinality (~7800 distinct, 5% db names)",
288+
&values,
289+
false,
290+
);
291+
}
292+
293+
criterion_group!(benches, array_agg_benchmark, distinct_array_agg_benchmark);
195294
criterion_main!(benches);

‎datafusion/functions-aggregate/src/array_agg.rs‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2589,4 +2589,128 @@ mod tests {
25892589

25902590
Ok(())
25912591
}
2592+
2593+
#[test]
2594+
fn distinct_array_agg_utf8_deduplicates() -> Result<()> {
2595+
use arrow::array::StringArray;
2596+
2597+
// 7 rows with 4 distinct values, each duplicate appearing twice.
2598+
let input: ArrayRef = Arc::new(StringArray::from(vec![
2599+
"postgres", "mysql", "postgres", "redis", "mysql", "duckdb", "redis",
2600+
]));
2601+
2602+
let mut acc = DistinctArrayAggAccumulator::try_new(&DataType::Utf8, None, false)?;
2603+
acc.update_batch(&[input])?;
2604+
2605+
let result = acc.evaluate()?;
2606+
let ScalarValue::List(arr) = &result else {
2607+
panic!("expected ScalarValue::List, got {result:?}");
2608+
};
2609+
2610+
let inner = arr.value(0);
2611+
let strings = inner
2612+
.as_any()
2613+
.downcast_ref::<StringArray>()
2614+
.expect("inner array should be StringArray");
2615+
2616+
// HashSet ordering is nondeterministic — sort before asserting.
2617+
let mut values: Vec<&str> =
2618+
(0..strings.len()).map(|i| strings.value(i)).collect();
2619+
values.sort_unstable();
2620+
2621+
assert_eq!(values, vec!["duckdb", "mysql", "postgres", "redis"]);
2622+
Ok(())
2623+
}
2624+
2625+
#[test]
2626+
fn distinct_array_agg_int64_deduplicates() -> Result<()> {
2627+
use arrow::array::Int64Array;
2628+
2629+
// 7 rows with 4 distinct values, each duplicate appearing twice.
2630+
let input: ArrayRef = Arc::new(Int64Array::from(vec![1i64, 2, 1, 3, 2, 4, 3]));
2631+
2632+
let mut acc =
2633+
DistinctArrayAggAccumulator::try_new(&DataType::Int64, None, false)?;
2634+
acc.update_batch(&[input])?;
2635+
2636+
let result = acc.evaluate()?;
2637+
let ScalarValue::List(arr) = &result else {
2638+
panic!("expected ScalarValue::List, got {result:?}");
2639+
};
2640+
2641+
let inner = arr.value(0);
2642+
let ints = inner
2643+
.as_any()
2644+
.downcast_ref::<Int64Array>()
2645+
.expect("inner array should be Int64Array");
2646+
2647+
let mut values: Vec<i64> = (0..ints.len()).map(|i| ints.value(i)).collect();
2648+
values.sort_unstable();
2649+
2650+
assert_eq!(values, vec![1i64, 2, 3, 4]);
2651+
Ok(())
2652+
}
2653+
2654+
#[test]
2655+
fn distinct_array_agg_float64_deduplicates() -> Result<()> {
2656+
use arrow::array::Float64Array;
2657+
2658+
// 7 rows with 4 distinct values, each duplicate appearing twice.
2659+
let input: ArrayRef = Arc::new(Float64Array::from(vec![
2660+
1.0f64, 2.5, 1.0, 3.75, 2.5, 4.0, 3.75,
2661+
]));
2662+
2663+
let mut acc =
2664+
DistinctArrayAggAccumulator::try_new(&DataType::Float64, None, false)?;
2665+
acc.update_batch(&[input])?;
2666+
2667+
let result = acc.evaluate()?;
2668+
let ScalarValue::List(arr) = &result else {
2669+
panic!("expected ScalarValue::List, got {result:?}");
2670+
};
2671+
2672+
let inner = arr.value(0);
2673+
let floats = inner
2674+
.as_any()
2675+
.downcast_ref::<Float64Array>()
2676+
.expect("inner array should be Float64Array");
2677+
2678+
// f64 has no Ord — use total_cmp for a stable sort.
2679+
let mut values: Vec<f64> = (0..floats.len()).map(|i| floats.value(i)).collect();
2680+
values.sort_unstable_by(|a, b| a.total_cmp(b));
2681+
2682+
assert_eq!(values, vec![1.0f64, 2.5, 3.75, 4.0]);
2683+
Ok(())
2684+
}
2685+
2686+
#[test]
2687+
fn distinct_array_agg_date32_deduplicates() -> Result<()> {
2688+
use arrow::array::Date32Array;
2689+
2690+
// 7 rows with 4 distinct dates (days since epoch), each duplicate appearing twice.
2691+
let input: ArrayRef = Arc::new(Date32Array::from(vec![
2692+
100i32, 200, 100, 300, 200, 400, 300,
2693+
]));
2694+
2695+
let mut acc =
2696+
DistinctArrayAggAccumulator::try_new(&DataType::Date32, None, false)?;
2697+
acc.update_batch(&[input])?;
2698+
2699+
let result = acc.evaluate()?;
2700+
let ScalarValue::List(arr) = &result else {
2701+
panic!("expected ScalarValue::List, got {result:?}");
2702+
};
2703+
2704+
let inner = arr.value(0);
2705+
let dates = inner
2706+
.as_any()
2707+
.downcast_ref::<Date32Array>()
2708+
.expect("inner array should be Date32Array");
2709+
2710+
let mut values: Vec<i32> = (0..dates.len()).map(|i| dates.value(i)).collect();
2711+
values.sort_unstable();
2712+
2713+
assert_eq!(values, vec![100i32, 200, 300, 400]);
2714+
Ok(())
2715+
}
25922716
}

0 commit comments

Comments
 (0)