From 20c43f162b8c4741350fcea73cf64f2ba83e1ce9 Mon Sep 17 00:00:00 2001 From: linfeng Date: Wed, 5 Aug 2026 21:09:23 +0800 Subject: [PATCH 1/2] Fix filtered aggregation for groups with no matching rows --- .../datafusion-ext-plans/src/agg/acc.rs | 19 +++++++++++ .../datafusion-ext-plans/src/agg/agg_ctx.rs | 4 +++ .../datafusion-ext-plans/src/agg_exec.rs | 32 +++++++++++-------- .../org/apache/auron/AuronQuerySuite.scala | 8 +++++ 4 files changed, 50 insertions(+), 13 deletions(-) diff --git a/native-engine/datafusion-ext-plans/src/agg/acc.rs b/native-engine/datafusion-ext-plans/src/agg/acc.rs index bea1c05ec..0857e161b 100644 --- a/native-engine/datafusion-ext-plans/src/agg/acc.rs +++ b/native-engine/datafusion-ext-plans/src/agg/acc.rs @@ -88,6 +88,25 @@ impl AccTable { self.cols.iter_mut().for_each(|c| c.resize(num_records)); } + pub fn ensure_size(&mut self, idx: IdxSelection<'_>) { + let num_records = match idx { + IdxSelection::Single(idx) => idx + 1, + IdxSelection::Indices(indices) => { + indices.iter().copied().max().map_or(0, |idx| idx + 1) + } + IdxSelection::IndicesU32(indices) => indices + .iter() + .copied() + .max() + .map_or(0, |idx| idx as usize + 1), + IdxSelection::Range(_, end) => end, + }; + self.cols + .iter_mut() + .filter(|col| col.num_records() < num_records) + .for_each(|col| col.resize(num_records)); + } + pub fn shrink_to_fit(&mut self) { self.cols.iter_mut().for_each(|c| c.shrink_to_fit()); } diff --git a/native-engine/datafusion-ext-plans/src/agg/agg_ctx.rs b/native-engine/datafusion-ext-plans/src/agg/agg_ctx.rs index f2dfbc240..73ce4a26b 100644 --- a/native-engine/datafusion-ext-plans/src/agg/agg_ctx.rs +++ b/native-engine/datafusion-ext-plans/src/agg/agg_ctx.rs @@ -274,6 +274,10 @@ impl AggContext { // arrow-ffi with sliced batch is buggy in older arrow-java, so we use unsliced // batch with explicit offsets + // Every group needs an accumulator slot even when FILTER excludes all of its + // rows. + acc_table.ensure_size(acc_idx); + // partial update if self.need_partial_update { let agg_exprs_batch = self.agg_expr_evaluator.filter_project(&batch)?; diff --git a/native-engine/datafusion-ext-plans/src/agg_exec.rs b/native-engine/datafusion-ext-plans/src/agg_exec.rs index d75d304f0..e5a573ec7 100644 --- a/native-engine/datafusion-ext-plans/src/agg_exec.rs +++ b/native-engine/datafusion-ext-plans/src/agg_exec.rs @@ -436,7 +436,7 @@ mod test { AggMode::{Final, Partial}, GroupingExpr, agg::create_agg, - sum::AggSum, + count::AggCount, }, agg_exec::AggExec, }; @@ -705,7 +705,9 @@ mod test { Field::new("flag", DataType::Boolean, false), ])); - let grp_col: ArrayRef = Arc::new(StringArray::from(vec!["a", "a", "a", "b", "b", "b"])); + let grp_col: ArrayRef = Arc::new(StringArray::from(vec![ + "a", "a", "a", "b", "b", "b", "c", "c", + ])); let val_col: ArrayRef = Arc::new(Int32Array::from(vec![ Some(1), Some(2), @@ -713,9 +715,11 @@ mod test { Some(4), Some(5), None, + Some(6), + Some(7), ])); let flag_col: ArrayRef = Arc::new(BooleanArray::from(vec![ - true, false, true, false, true, true, + true, false, true, false, true, true, false, false, ])); let batch = RecordBatch::try_new( @@ -729,8 +733,9 @@ mod test { None, )?); - // SUM(val) FILTER (WHERE flag = true) - // Expected: a=1+3=4, b=5+0=5 (NULL val contributes 0 to sum) + // COUNT(val) FILTER (WHERE flag = true) + // Expected: a=2, b=1 (NULL val is not counted), c=0 because all rows in + // group c are filtered out. let filter_expr = Arc::new(phys_expr::Column::new("flag", 2)); let agg_exec = Arc::new(AggExec::try_new( HashAgg, @@ -739,11 +744,11 @@ mod test { expr: phys_expr::col("grp", &schema)?, }], vec![AggExpr { - field_name: "sum_filtered".to_string(), + field_name: "count_filtered".to_string(), mode: Partial, filter: Some(filter_expr), - agg: Arc::new(AggSum::try_new( - phys_expr::col("val", &schema)?, + agg: Arc::new(AggCount::try_new( + vec![phys_expr::col("val", &schema)?], DataType::Int64, )?), }], @@ -755,15 +760,16 @@ mod test { let result = concat_batches(&output[0].schema(), &output)?; let grp_result = result.column(0).as_string::(); - let sum_result = result.column(1).as_primitive::(); + let count_result = result.column(1).as_primitive::(); - assert_eq!(grp_result.len(), 2); + assert_eq!(grp_result.len(), 3); let mut found = std::collections::HashMap::new(); for i in 0..grp_result.len() { - found.insert(grp_result.value(i), sum_result.value(i)); + found.insert(grp_result.value(i), count_result.value(i)); } - assert_eq!(found["a"], 4); // 1 + 3 - assert_eq!(found["b"], 5); // 5 (NULL val contributes 0) + assert_eq!(found["a"], 2); + assert_eq!(found["b"], 1); + assert_eq!(found["c"], 0); Ok(()) } diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala index 56119a7cb..e080baaa6 100644 --- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala +++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala @@ -974,6 +974,14 @@ class AuronQuerySuite extends AuronQueryTest with BaseAuronSQLSuite with AuronSQ |GROUP BY category |ORDER BY category""".stripMargin) + // A group with no matching rows keeps the aggregate's empty-input state. + checkSparkAnswerAndOperator("""SELECT category, + | SUM(amount) FILTER (WHERE amount > 450) AS high_amount, + | SUM(amount) FILTER (WHERE amount < 150) AS low_amount + |FROM t_filter_agg_2289 + |GROUP BY category + |ORDER BY category""".stripMargin) + // Multiple aggregates with different FILTER predicates checkSparkAnswerAndOperator("""SELECT | SUM(amount) FILTER (WHERE is_vip = true) AS sum_vip, From ce31ef224487dd8edb27d182fc49a67be7c48952 Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Thu, 6 Aug 2026 10:41:00 +0800 Subject: [PATCH 2/2] add filtered count regression test --- .../datafusion-ext-plans/src/agg_exec.rs | 58 +++++++++++-------- 1 file changed, 34 insertions(+), 24 deletions(-) diff --git a/native-engine/datafusion-ext-plans/src/agg_exec.rs b/native-engine/datafusion-ext-plans/src/agg_exec.rs index e5a573ec7..7ffaf295f 100644 --- a/native-engine/datafusion-ext-plans/src/agg_exec.rs +++ b/native-engine/datafusion-ext-plans/src/agg_exec.rs @@ -437,6 +437,7 @@ mod test { GroupingExpr, agg::create_agg, count::AggCount, + sum::AggSum, }, agg_exec::AggExec, }; @@ -705,9 +706,7 @@ mod test { Field::new("flag", DataType::Boolean, false), ])); - let grp_col: ArrayRef = Arc::new(StringArray::from(vec![ - "a", "a", "a", "b", "b", "b", "c", "c", - ])); + let grp_col: ArrayRef = Arc::new(StringArray::from(vec!["a", "a", "a", "b", "b", "b"])); let val_col: ArrayRef = Arc::new(Int32Array::from(vec![ Some(1), Some(2), @@ -715,11 +714,9 @@ mod test { Some(4), Some(5), None, - Some(6), - Some(7), ])); let flag_col: ArrayRef = Arc::new(BooleanArray::from(vec![ - true, false, true, false, true, true, false, false, + true, false, true, false, true, true, ])); let batch = RecordBatch::try_new( @@ -733,9 +730,8 @@ mod test { None, )?); - // COUNT(val) FILTER (WHERE flag = true) - // Expected: a=2, b=1 (NULL val is not counted), c=0 because all rows in - // group c are filtered out. + // SUM(val) FILTER (WHERE flag = true) + // Expected: a=1+3=4, b=5+0=5 (NULL val contributes 0 to sum) let filter_expr = Arc::new(phys_expr::Column::new("flag", 2)); let agg_exec = Arc::new(AggExec::try_new( HashAgg, @@ -743,15 +739,28 @@ mod test { field_name: "grp".to_string(), expr: phys_expr::col("grp", &schema)?, }], - vec![AggExpr { - field_name: "count_filtered".to_string(), - mode: Partial, - filter: Some(filter_expr), - agg: Arc::new(AggCount::try_new( - vec![phys_expr::col("val", &schema)?], - DataType::Int64, - )?), - }], + vec![ + AggExpr { + field_name: "sum_filtered".to_string(), + mode: Partial, + filter: Some(filter_expr), + agg: Arc::new(AggSum::try_new( + phys_expr::col("val", &schema)?, + DataType::Int64, + )?), + }, + AggExpr { + field_name: "count_filtered".to_string(), + mode: Partial, + filter: Some(Arc::new(phys_expr::Literal::new(ScalarValue::Boolean( + Some(false), + )))), + agg: Arc::new(AggCount::try_new( + vec![phys_expr::col("val", &schema)?], + DataType::Int64, + )?), + }, + ], false, input, )?); @@ -760,16 +769,17 @@ mod test { let result = concat_batches(&output[0].schema(), &output)?; let grp_result = result.column(0).as_string::(); - let count_result = result.column(1).as_primitive::(); + let sum_result = result.column(1).as_primitive::(); + let count_result = result.column(2).as_primitive::(); - assert_eq!(grp_result.len(), 3); + assert_eq!(grp_result.len(), 2); let mut found = std::collections::HashMap::new(); for i in 0..grp_result.len() { - found.insert(grp_result.value(i), count_result.value(i)); + found.insert(grp_result.value(i), sum_result.value(i)); + assert_eq!(count_result.value(i), 0); } - assert_eq!(found["a"], 2); - assert_eq!(found["b"], 1); - assert_eq!(found["c"], 0); + assert_eq!(found["a"], 4); // 1 + 3 + assert_eq!(found["b"], 5); // 5 (NULL val contributes 0) Ok(()) }