fix[physical-plan/aggregates]: fix grouping by Ree<Dict>

datafusion/core/tests/sql/aggregates/basic.rs

@@ -441,3 +441,58 @@ async fn count_distinct_dictionary_mixed_values() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn group_by_ree_dict_column() -> Result<()> {
+    let ctx = SessionContext::new();
+
+    let run_ends = Int32Array::from(vec![2, 4, 5]);
+    let dict = DictionaryArray::new(
+        UInt32Array::from(vec![0, 1, 2]),
+        Arc::new(StringArray::from(vec!["alpha", "beta", "gamma"])),
+    );
+    let ree_col = RunArray::<Int32Type>::try_new(&run_ends, &dict).unwrap();
+    let value_col = Int32Array::from(vec![1, 2, 3, 4, 5]);
+
+    let dict_type =
+        DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8));
+    let schema = Arc::new(Schema::new(vec![
+        Field::new(
+            "group_col",
+            DataType::RunEndEncoded(
+                Arc::new(Field::new("run_ends", DataType::Int32, false)),
+                Arc::new(Field::new("values", dict_type, true)),
+            ),
+            true,
+        ),
+        Field::new("value", DataType::Int32, false),
+    ]));
+
+    let batch = RecordBatch::try_new(
+        schema.clone(),
+        vec![Arc::new(ree_col), Arc::new(value_col)],
+    )?;
+    let table = MemTable::try_new(schema, vec![vec![batch]])?;
+    ctx.register_table("t", Arc::new(table))?;
+
+    let results = ctx
+        .sql("SELECT group_col, SUM(value) as total FROM t GROUP BY group_col ORDER BY group_col")
+        .await?
+        .collect()
+        .await?;
+
+    assert_snapshot!(
+        batches_to_string(&results),
+        @r"
+    +-----------+-------+
+    | group_col | total |
+    +-----------+-------+
+    | alpha     | 3     |
+    | beta      | 7     |
+    | gamma     | 5     |
+    +-----------+-------+
+    "
+    );
+
+    Ok(())
+}

datafusion/physical-plan/src/aggregates/group_values/row.rs

@@ -16,7 +16,10 @@
 // under the License.
 
 use crate::aggregates::group_values::GroupValues;
-use arrow::array::{Array, ArrayRef, ListArray, StructArray};
+use arrow::array::{
+    Array, ArrayRef, ListArray, PrimitiveArray, RunArray, StructArray,
+    downcast_run_end_index,
+};
 use arrow::compute::cast;
 use arrow::datatypes::{DataType, SchemaRef};
 use arrow::row::{RowConverter, Rows, SortField};
@@ -291,6 +294,33 @@ fn dictionary_encode_if_necessary(
             )?))
         }
         (DataType::Dictionary(_, _), _) => Ok(cast(array.as_ref(), expected)?),
+        (
+            DataType::RunEndEncoded(run_ends_field, expected_values_field),
+            &DataType::RunEndEncoded(_, _),
+        ) => {
+            macro_rules! reencode_ree {
+                ($run_end_type:ty) => {{
+                    let run_array = array
+                        .as_any()
+                        .downcast_ref::<RunArray<$run_end_type>>()
+                        .unwrap();
+                    let values = dictionary_encode_if_necessary(
+                        &(Arc::clone(run_array.values()) as ArrayRef),
+                        expected_values_field.data_type(),
+                    )?;
+                    let run_ends = PrimitiveArray::<$run_end_type>::new(
+                        run_array.run_ends().inner().clone(),
+                        None,
+                    );
+                    Ok(Arc::new(RunArray::try_new(&run_ends, &values)?))
+                }};
+            }
+            downcast_run_end_index! {
+                run_ends_field.data_type() => (reencode_ree),
+                _ => unreachable!("unsupported run end type: {}", run_ends_field.data_type()),
+            }
+        }
+        (DataType::RunEndEncoded(_, _), _) => Ok(cast(array.as_ref(), expected)?),
         (_, _) => Ok(Arc::<dyn Array>::clone(array)),
     }
 }