You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/04/27 17:10:25 UTC
(datafusion) branch gh-readonly-queue/main/pr-21548-40b209ee6700723d1179e98b5862b7e5a7e2b6a5 created (now ca1d39dfca)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-21548-40b209ee6700723d1179e98b5862b7e5a7e2b6a5
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at ca1d39dfca perf: implement convert_to_state for SparkAvg (#21548)
This branch includes the following new commits:
new ca1d39dfca perf: implement convert_to_state for SparkAvg (#21548)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: perf: implement convert_to_state for SparkAvg (#21548)
Posted by gi...@apache.org.
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-21548-40b209ee6700723d1179e98b5862b7e5a7e2b6a5
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit ca1d39dfca531620b69a3a17018a899bb478a361
Author: Alex Zhang <87...@users.noreply.github.com>
AuthorDate: Mon Apr 27 10:10:03 2026 -0700
perf: implement convert_to_state for SparkAvg (#21548)
## Which issue does this PR close?
- Part of #17964.
## Rationale for this change
SparkAvg's AvgGroupsAccumulator doesn't implement
supports_convert_to_state (defaults to false), which prevents the
skip-partial-aggregation optimization from kicking in for queries that
use Spark's avg().
I ran into this while benchmarking a Spark Connect engine built on
DataFusion. On TPC-H q17 at SF10, the partial aggregate for
avg(l_quantity) grouped by l_partkey (~2M groups out of 60M rows) was
not triggering skip-aggregation:
| Metric | Without convert_to_state | With convert_to_state |
|--------|-------------------------|-----------------------|
| Partial aggregate memory | 923 MB | 40 MB |
| Partial aggregate elapsed | 4.75s | 109ms |
The skip-aggregation probe (#11627) detects when a partial aggregate
isn't reducing cardinality and falls back to passing rows through as
state directly. This needs convert_to_state so the accumulator can
produce [sum, count] state arrays from raw input. The built-in Avg
already has this (#11734), but it wasn't carried over when SparkAvg was
migrated from Comet in #17871.
## What changes are included in this PR?
Adds convert_to_state() and supports_convert_to_state() to
AvgGroupsAccumulator in datafusion-spark.
Follows the same approach as the built-in Avg, adapted for SparkAvg's
differences:
- State order is [sum, count] (vs [count, sum] in the built-in)
- Count type is Int64 (vs UInt64 in the built-in)
- Null handling uses NullBuffer::union directly instead of pulling in
datafusion-functions-aggregate-common as a dep
Also cleaned up the fully-qualified arrow::array::BooleanArray
references in update_batch / merge_batch since adding BooleanArray to
the import block triggered the unused_qualifications lint.
## Are these changes tested?
Yes, unit tests covering basic conversion, null propagation, filter
handling, and a roundtrip through merge_batch to verify the converted
state produces correct results end-to-end.
## Are there any user-facing changes?
No. Queries using avg() through the Spark function registry will
automatically benefit from skip-partial-aggregation on high-cardinality
groupings.
---------
Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
---
Cargo.lock | 1 +
datafusion/spark/Cargo.toml | 1 +
datafusion/spark/src/function/aggregate/avg.rs | 172 +++++++++++++++++++++++--
3 files changed, 161 insertions(+), 13 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index ed3c4825bd..9d60540d9c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2632,6 +2632,7 @@ dependencies = [
"datafusion-expr",
"datafusion-functions",
"datafusion-functions-aggregate",
+ "datafusion-functions-aggregate-common",
"datafusion-functions-nested",
"log",
"num-traits",
diff --git a/datafusion/spark/Cargo.toml b/datafusion/spark/Cargo.toml
index 0651fe5651..72bca58b7a 100644
--- a/datafusion/spark/Cargo.toml
+++ b/datafusion/spark/Cargo.toml
@@ -55,6 +55,7 @@ datafusion-execution = { workspace = true }
datafusion-expr = { workspace = true }
datafusion-functions = { workspace = true, features = ["crypto_expressions"] }
datafusion-functions-aggregate = { workspace = true }
+datafusion-functions-aggregate-common = { workspace = true }
datafusion-functions-nested = { workspace = true }
log = { workspace = true }
num-traits = { workspace = true }
diff --git a/datafusion/spark/src/function/aggregate/avg.rs b/datafusion/spark/src/function/aggregate/avg.rs
index c504f5387f..5f4d2c253a 100644
--- a/datafusion/spark/src/function/aggregate/avg.rs
+++ b/datafusion/spark/src/function/aggregate/avg.rs
@@ -16,7 +16,8 @@
// under the License.
use arrow::array::{
- Array, ArrayRef, ArrowNativeTypeOp, ArrowNumericType, Int64Array, PrimitiveArray,
+ Array, ArrayRef, ArrowNativeTypeOp, ArrowNumericType, BooleanArray, Int64Array,
+ PrimitiveArray,
builder::PrimitiveBuilder,
cast::AsArray,
types::{Float64Type, Int64Type},
@@ -31,6 +32,9 @@ use datafusion_expr::{
Accumulator, AggregateUDFImpl, Coercion, EmitTo, GroupsAccumulator, ReversedUDAF,
Signature, TypeSignatureClass, Volatility,
};
+use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::{
+ filtered_null_mask, set_nulls,
+};
use std::sync::Arc;
/// AVG aggregate expression
@@ -248,7 +252,7 @@ where
&mut self,
values: &[ArrayRef],
group_indices: &[usize],
- _opt_filter: Option<&arrow::array::BooleanArray>,
+ _opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> Result<()> {
assert_eq!(values.len(), 1, "single argument to update_batch");
@@ -285,26 +289,26 @@ where
&mut self,
values: &[ArrayRef],
group_indices: &[usize],
- _opt_filter: Option<&arrow::array::BooleanArray>,
+ _opt_filter: Option<&BooleanArray>,
total_num_groups: usize,
) -> Result<()> {
assert_eq!(values.len(), 2, "two arguments to merge_batch");
// first batch is partial sums, second is counts
let partial_sums = values[0].as_primitive::<T>();
let partial_counts = values[1].as_primitive::<Int64Type>();
- // update counts with partial counts
- self.counts.resize(total_num_groups, 0);
- let iter1 = group_indices.iter().zip(partial_counts.values().iter());
- for (&group_index, &partial_count) in iter1 {
- self.counts[group_index] += partial_count;
- }
- // update sums
+ self.counts.resize(total_num_groups, 0);
self.sums.resize(total_num_groups, T::default_value());
- let iter2 = group_indices.iter().zip(partial_sums.values().iter());
- for (&group_index, &new_value) in iter2 {
+
+ for (idx, &group_index) in group_indices.iter().enumerate() {
+ // Skip null state entries emitted by convert_to_state for
+ // filtered / null input rows.
+ if partial_counts.is_null(idx) || partial_sums.is_null(idx) {
+ continue;
+ }
+ self.counts[group_index] += partial_counts.value(idx);
let sum = &mut self.sums[group_index];
- *sum = sum.add_wrapping(new_value);
+ *sum = sum.add_wrapping(partial_sums.value(idx));
}
Ok(())
@@ -343,7 +347,149 @@ where
])
}
+ fn convert_to_state(
+ &self,
+ values: &[ArrayRef],
+ opt_filter: Option<&BooleanArray>,
+ ) -> Result<Vec<ArrayRef>> {
+ let sums = values[0]
+ .as_primitive::<T>()
+ .clone()
+ .with_data_type(self.return_data_type.clone());
+ let counts = Int64Array::from_value(1, sums.len());
+
+ let nulls = filtered_null_mask(opt_filter, &sums);
+ let counts = set_nulls(counts, nulls.clone());
+ let sums = set_nulls(sums, nulls);
+
+ // [sum, count] - must match state() and merge_batch()
+ Ok(vec![
+ Arc::new(sums) as ArrayRef,
+ Arc::new(counts) as ArrayRef,
+ ])
+ }
+
+ fn supports_convert_to_state(&self) -> bool {
+ true
+ }
+
fn size(&self) -> usize {
self.counts.capacity() * size_of::<i64>() + self.sums.capacity() * size_of::<T>()
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::Float64Array;
+
+ fn make_acc() -> AvgGroupsAccumulator<Float64Type, impl Fn(f64, i64) -> Result<f64>> {
+ AvgGroupsAccumulator::<Float64Type, _>::new(&DataType::Float64, |sum, count| {
+ Ok(sum / count as f64)
+ })
+ }
+
+ #[test]
+ fn supports_convert_to_state() {
+ assert!(make_acc().supports_convert_to_state());
+ }
+
+ #[test]
+ fn convert_to_state_basic() {
+ let acc = make_acc();
+ let values: Vec<ArrayRef> =
+ vec![Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0]))];
+ let state = acc.convert_to_state(&values, None).unwrap();
+
+ assert_eq!(state.len(), 2);
+ let sums = state[0].as_primitive::<Float64Type>();
+ let counts = state[1].as_primitive::<Int64Type>();
+
+ assert_eq!(sums.values().as_ref(), &[1.0, 2.0, 3.0]);
+ assert_eq!(counts.values().as_ref(), &[1, 1, 1]);
+ assert_eq!(sums.null_count(), 0);
+ assert_eq!(counts.null_count(), 0);
+ }
+
+ #[test]
+ fn convert_to_state_with_nulls() {
+ let acc = make_acc();
+ let values: Vec<ArrayRef> = vec![Arc::new(Float64Array::from(vec![
+ Some(1.0),
+ None,
+ Some(3.0),
+ ]))];
+ let state = acc.convert_to_state(&values, None).unwrap();
+
+ let sums = state[0].as_primitive::<Float64Type>();
+ let counts = state[1].as_primitive::<Int64Type>();
+
+ assert!(!sums.is_null(0));
+ assert!(sums.is_null(1));
+ assert!(!sums.is_null(2));
+
+ assert_eq!(counts.value(0), 1);
+ assert!(counts.is_null(1));
+ assert_eq!(counts.value(2), 1);
+ }
+
+ #[test]
+ fn convert_to_state_with_filter() {
+ let acc = make_acc();
+ let values: Vec<ArrayRef> =
+ vec![Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0]))];
+ let filter = BooleanArray::from(vec![true, false, true]);
+ let state = acc.convert_to_state(&values, Some(&filter)).unwrap();
+
+ let sums = state[0].as_primitive::<Float64Type>();
+ let counts = state[1].as_primitive::<Int64Type>();
+
+ assert!(!sums.is_null(0));
+ assert!(sums.is_null(1));
+ assert!(!sums.is_null(2));
+
+ assert_eq!(counts.value(0), 1);
+ assert!(counts.is_null(1));
+ assert_eq!(counts.value(2), 1);
+ }
+
+ #[test]
+ fn convert_to_state_roundtrips_through_merge() {
+ let mut acc = make_acc();
+ let input: Vec<ArrayRef> =
+ vec![Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0]))];
+ let state = acc.convert_to_state(&input, None).unwrap();
+
+ // feed the converted state back through merge_batch
+ acc.merge_batch(
+ &state,
+ &[0, 0, 0],
+ None,
+ 1, // single group
+ )
+ .unwrap();
+
+ let result = acc.evaluate(EmitTo::All).unwrap();
+ let result = result.as_primitive::<Float64Type>();
+ assert_eq!(result.value(0), 20.0); // (10+20+30)/3
+ }
+
+ #[test]
+ fn convert_to_state_null_merge_matches_direct() {
+ // avg([1.0, NULL, 3.0]) must be 2.0 after a convert_to_state → merge_batch
+ // round-trip. Before the merge-path null fix this leaked the backing
+ // buffer value at the null slot and produced the wrong average.
+ let mut acc = make_acc();
+ let input: Vec<ArrayRef> = vec![Arc::new(Float64Array::from(vec![
+ Some(1.0),
+ None,
+ Some(3.0),
+ ]))];
+ let state = acc.convert_to_state(&input, None).unwrap();
+ acc.merge_batch(&state, &[0, 0, 0], None, 1).unwrap();
+
+ let result = acc.evaluate(EmitTo::All).unwrap();
+ let result = result.as_primitive::<Float64Type>();
+ assert_eq!(result.value(0), 2.0);
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org