Fix convert_to_state bug in GroupsAccumulatorAdapter (#12834)

* Fix convert_to_state bug * fmt
apache · Oct 9, 2024 · 0c7eb0d · 0c7eb0d
1 parent ffe73fe
commit 0c7eb0d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs
@@ -254,10 +254,8 @@ async fn test_basic_string_aggr_group_by_single_int64() {
     let fuzzer = builder
         .data_gen_config(data_gen_config)
         .data_gen_rounds(8)
-        // FIXME: Encounter error in min/max
-        //   ArrowError(InvalidArgumentError("number of columns(1) must match number of fields(2) in schema"))
-        // .add_sql("SELECT b, max(a) FROM fuzz_table GROUP BY b")
-        // .add_sql("SELECT b, min(a) FROM fuzz_table GROUP BY b")
+        .add_sql("SELECT b, max(a) FROM fuzz_table GROUP BY b")
+        .add_sql("SELECT b, min(a) FROM fuzz_table GROUP BY b")
         .add_sql("SELECT b, count(a) FROM fuzz_table GROUP BY b")
         .add_sql("SELECT b, count(distinct a) FROM fuzz_table GROUP BY b")
         .table_name("fuzz_table")
@@ -291,10 +289,8 @@ async fn test_basic_string_aggr_group_by_single_string() {
     let fuzzer = builder
         .data_gen_config(data_gen_config)
         .data_gen_rounds(16)
-        // FIXME: Encounter error in min/max
-        //   ArrowError(InvalidArgumentError("number of columns(1) must match number of fields(2) in schema"))
-        // .add_sql("SELECT b, max(a) FROM fuzz_table GROUP BY b")
-        // .add_sql("SELECT b, min(a) FROM fuzz_table GROUP BY b")
+        .add_sql("SELECT b, max(a) FROM fuzz_table GROUP BY b")
+        .add_sql("SELECT b, min(a) FROM fuzz_table GROUP BY b")
         .add_sql("SELECT b, count(a) FROM fuzz_table GROUP BY b")
         .add_sql("SELECT b, count(distinct a) FROM fuzz_table GROUP BY b")
         .table_name("fuzz_table")
@@ -329,10 +325,8 @@ async fn test_basic_string_aggr_group_by_mixed_string_int64() {
     let fuzzer = builder
         .data_gen_config(data_gen_config)
         .data_gen_rounds(16)
-        // FIXME: Encounter error in min/max
-        //   ArrowError(InvalidArgumentError("number of columns(1) must match number of fields(2) in schema"))
-        // .add_sql("SELECT b, c, max(a) FROM fuzz_table GROUP BY b, c")
-        // .add_sql("SELECT b, c, min(a) FROM fuzz_table GROUP BY b, c")
+        .add_sql("SELECT b, c, max(a) FROM fuzz_table GROUP BY b, c")
+        .add_sql("SELECT b, c, min(a) FROM fuzz_table GROUP BY b, c")
         .add_sql("SELECT b, c, count(a) FROM fuzz_table GROUP BY b, c")
         .add_sql("SELECT b, c, count(distinct a) FROM fuzz_table GROUP BY b, c")
         .table_name("fuzz_table")

diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs
@@ -23,6 +23,7 @@ pub mod bool_op;
 pub mod nulls;
 pub mod prim_op;
 
+use arrow::array::new_empty_array;
 use arrow::{
     array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray},
     compute,
@@ -405,6 +406,18 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
     ) -> Result<Vec<ArrayRef>> {
         let num_rows = values[0].len();
 
+        // If there are no rows, return empty arrays
+        if num_rows == 0 {
+            // create empty accumulator to get the state types
+            let empty_state = (self.factory)()?.state()?;
+            let empty_arrays = empty_state
+                .into_iter()
+                .map(|state_val| new_empty_array(&state_val.data_type()))
+                .collect::<Vec<_>>();
+
+            return Ok(empty_arrays);
+        }
+
         // Each row has its respective group
         let mut results = vec![];
         for row_idx in 0..num_rows {