[AURON #2453] Fix filtered aggregation for groups with no matching rows

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());
     }

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)?;

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::<i32>();
         let sum_result = result.column(1).as_primitive::<Int64Type>();
+        let count_result = result.column(2).as_primitive::<Int64Type>();
 
         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)

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,