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..7ffaf295f 100644 --- a/native-engine/datafusion-ext-plans/src/agg_exec.rs +++ b/native-engine/datafusion-ext-plans/src/agg_exec.rs @@ -436,6 +436,7 @@ mod test { AggMode::{Final, Partial}, GroupingExpr, agg::create_agg, + count::AggCount, sum::AggSum, }, agg_exec::AggExec, @@ -738,15 +739,28 @@ mod test { field_name: "grp".to_string(), expr: phys_expr::col("grp", &schema)?, }], - 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, - )?), - }], + 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, )?); @@ -756,11 +770,13 @@ mod test { let grp_result = result.column(0).as_string::(); let sum_result = result.column(1).as_primitive::(); + let count_result = result.column(2).as_primitive::(); 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), sum_result.value(i)); + assert_eq!(count_result.value(i), 0); } assert_eq!(found["a"], 4); // 1 + 3 assert_eq!(found["b"], 5); // 5 (NULL val contributes 0) 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,