Feat: add dictionaries as a supported group column type

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs

@@ -0,0 +1,373 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::aggregates::group_values::multi_group_by::GroupColumn;
+use arrow::array::{
+    Array, ArrayRef, AsArray, BooleanBufferBuilder, DictionaryArray, PrimitiveArray,
+};
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType, Field};
+use datafusion_common::{Result, exec_err};
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+/// A [`GroupColumn`] implementation for dictionary arrays.
+/// wraps an inner [`GroupColumn`] that stores the dictionary values, and uses a null sentinel for null keys.
+pub struct DictionaryGroupValuesColumn<K: ArrowDictionaryKeyType + Send + Sync> {
+    inner: Box<dyn GroupColumn>,
+    // unary array used as a null sentinel for dictionary keys that are null
+    null_array: ArrayRef,
+    _phantom: PhantomData<K>,
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> {
+    pub fn new(inner: Box<dyn GroupColumn>, field: &Field) -> Self {
+        let null_array = arrow::array::new_null_array(field.data_type(), 1);
+        Self {
+            inner,
+            null_array,
+            _phantom: PhantomData,
+        }
+    }
+
+    fn into_dict(values: ArrayRef) -> ArrayRef {
+        let num_values = values.len();
+        let keys: PrimitiveArray<K> = (0..num_values)
+            .map(|i| (!values.is_null(i)).then(|| K::Native::usize_as(i)))
+            .collect();
+        Arc::new(DictionaryArray::<K>::new(keys, values))
+    }
+
+    /// Returns the index of the first null in `values`, or `None` if there are no nulls.
+    fn null_sentinel_index(values: &ArrayRef) -> Option<usize> {
+        values.nulls()?.iter().position(|valid| !valid)
+    }
+
+    /// Returns an error if adding `additional` groups would overflow the key type's range.
+    // https://github.com/apache/datafusion/issues/23127
+    fn check_key_overflow(current_len: usize, additional: usize) -> Result<()> {
+        if !Self::valid_bounds::<K>(current_len + additional) {
+            return exec_err!(
+                "Dictionary key type {:?} cannot represent {} groups",
+                K::DATA_TYPE,
+                current_len + additional
+            );
+        }
+        Ok(())
+    }
+
+    fn valid_bounds<T: ArrowDictionaryKeyType>(num_values: usize) -> bool {
+        let max: usize = match T::DATA_TYPE {
+            DataType::Int8 => i8::MAX as usize,
+            DataType::Int16 => i16::MAX as usize,
+            DataType::Int32 => i32::MAX as usize,
+            DataType::Int64 => i64::MAX as usize,
+            DataType::UInt8 => u8::MAX as usize,
+            DataType::UInt16 => u16::MAX as usize,
+            DataType::UInt32 => u32::MAX as usize,
+            DataType::UInt64 => usize::MAX,
+            _ => return false,
+        };
+        num_values == 0 || num_values - 1 <= max
+    }
+}
+
+impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
+    for DictionaryGroupValuesColumn<K>
+{
+    fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool {
+        let dict = array.as_dictionary::<K>();
+        match dict.key(rhs_row) {
+            None => self.inner.equal_to(lhs_row, &self.null_array, 0),
+            Some(key) => self.inner.equal_to(lhs_row, dict.values(), key),
+        }
+    }
+
+    fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> {
+        Self::check_key_overflow(self.inner.len(), 1)?;
+        let dict = array.as_dictionary::<K>();
+        match dict.key(row) {
+            None => self.inner.append_val(&self.null_array, 0),
+            Some(key) => self.inner.append_val(dict.values(), key),
+        }
+    }
+
+    fn vectorized_equal_to(
+        &self,
+        lhs_rows: &[usize],
+        array: &ArrayRef,
+        rhs_rows: &[usize],
+        equal_to_results: &mut BooleanBufferBuilder,
+    ) {
+        let dict = array.as_dictionary::<K>();
+        let dict_values = dict.values();
+        let dict_keys = dict.keys();
+
+        if dict_keys.null_count() == 0 {
+            // Fast path: no nulls in the key array, resolve all indices up front and
+            // delegate to the inner column's vectorized comparison.
+            let value_indices: Vec<usize> = rhs_rows
+                .iter()
+                .map(|&row_index| dict_keys.value(row_index).as_usize())
+                .collect();
+            self.inner.vectorized_equal_to(
+                lhs_rows,
+                dict_values,
+                &value_indices,
+                equal_to_results,
+            );
+        } else if let Some(sentinel) = Self::null_sentinel_index(dict_values) {
+            // Values already contains a null entry; reuse its position as the
+            // sentinel so null keys map there without any allocation.
+            let value_indices: Vec<usize> = rhs_rows
+                .iter()
+                .map(|&row_index| dict.key(row_index).unwrap_or(sentinel))
+                .collect();
+            self.inner.vectorized_equal_to(
+                lhs_rows,
+                dict_values,
+                &value_indices,
+                equal_to_results,
+            );
+        } else {
+            // Null keys but values has no null entry: scalar fallback per row.
+            for (i, (&lhs_row, &rhs_row)) in
+                lhs_rows.iter().zip(rhs_rows.iter()).enumerate()
+            {
+                if equal_to_results.get_bit(i) {
+                    let eq = match dict.key(rhs_row) {
+                        None => self.inner.equal_to(lhs_row, &self.null_array, 0),
+                        Some(key) => self.inner.equal_to(lhs_row, dict_values, key),
+                    };
+                    if !eq {
+                        equal_to_results.set_bit(i, false);
+                    }
+                }
+            }
+        }
+    }
+
+    fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> {
+        Self::check_key_overflow(self.inner.len(), rows.len())?;
+        let dict = array.as_dictionary::<K>();
+        let dict_keys = dict.keys();
+
+        if dict_keys.null_count() == 0 {
+            let value_indices: Vec<usize> = rows
+                .iter()
+                .map(|&row_index| dict_keys.value(row_index).as_usize())
+                .collect();
+            self.inner.vectorized_append(dict.values(), &value_indices)
+        } else if let Some(sentinel) = Self::null_sentinel_index(dict.values()) {
+            // Values already contains a null entry; reuse its position as the
+            // sentinel so null keys map there without any allocation.
+            let value_indices: Vec<usize> = rows
+                .iter()
+                .map(|&row_index| dict.key(row_index).unwrap_or(sentinel))
+                .collect();
+            self.inner.vectorized_append(dict.values(), &value_indices)
+        } else {
+            // Null keys but values has no null entry: scalar fallback per row.
+            for &row_index in rows {
+                match dict.key(row_index) {
+                    None => self.inner.append_val(&self.null_array, 0)?,
+                    Some(key) => self.inner.append_val(dict.values(), key)?,
+                }
+            }
+            Ok(())
+        }
+    }
+
+    fn len(&self) -> usize {
+        self.inner.len()
+    }
+
+    fn size(&self) -> usize {
+        self.inner.size()
+    }
+
+    fn build(self: Box<Self>) -> ArrayRef {
+        Self::into_dict(self.inner.build())
+    }
+
+    fn take_n(&mut self, n: usize) -> ArrayRef {
+        Self::into_dict(self.inner.take_n(n))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::aggregates::group_values::multi_group_by::bytes::ByteGroupValueBuilder;
+    use arrow::array::{
+        Array, ArrayRef, BooleanBufferBuilder, DictionaryArray, Int32Array, StringArray,
+    };
+    use arrow::compute::cast;
+    use arrow::datatypes::{DataType, Int32Type};
+    use datafusion_physical_expr::binary_map::OutputType;
+    use std::sync::Arc;
+
+    fn utf8_col() -> DictionaryGroupValuesColumn<Int32Type> {
+        let field = Field::new("", DataType::Utf8, true);
+        DictionaryGroupValuesColumn::<Int32Type>::new(
+            Box::new(ByteGroupValueBuilder::<i32>::new(OutputType::Utf8)),
+            &field,
+        )
+    }
+
+    fn dict_arr(keys: &[Option<i32>], values: &[&str]) -> ArrayRef {
+        Arc::new(DictionaryArray::<Int32Type>::new(
+            Int32Array::from(keys.to_vec()),
+            Arc::new(StringArray::from(values.to_vec())),
+        ))
+    }
+
+    fn str_values(arr: &ArrayRef) -> Vec<Option<String>> {
+        let plain = cast(arr.as_ref(), &DataType::Utf8).unwrap();
+        let strings = plain.as_any().downcast_ref::<StringArray>().unwrap();
+        (0..strings.len())
+            .map(|i| strings.is_valid(i).then(|| strings.value(i).to_owned()))
+            .collect()
+    }
+
+    fn assert_is_dict_utf8(arr: &ArrayRef) {
+        assert!(
+            matches!(arr.data_type(),
+                DataType::Dictionary(k, v)
+                    if k.as_ref() == &DataType::Int32 && v.as_ref() == &DataType::Utf8),
+            "expected Dictionary(Int32, Utf8), got {:?}",
+            arr.data_type()
+        );
+    }
+
+    fn true_buf(len: usize) -> BooleanBufferBuilder {
+        let mut buf = BooleanBufferBuilder::new(len);
+        buf.append_n(len, true);
+        buf
+    }
+
+    fn buf_to_vec(buf: &BooleanBufferBuilder) -> Vec<bool> {
+        (0..buf.len()).map(|i| buf.get_bit(i)).collect()
+    }
+
+    #[test]
+    fn utf8_dict_null_handling() {
+        let mut col = utf8_col();
+        let input: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::new(
+            Int32Array::from(vec![None, Some(0), Some(1)]),
+            Arc::new(StringArray::from(vec![None::<&str>, Some("b")])),
+        ));
+        for row in 0..3 {
+            col.append_val(&input, row).unwrap();
+        }
+
+        assert!(col.equal_to(0, &input, 0) && col.equal_to(1, &input, 1));
+        assert!(col.equal_to(0, &input, 1)); // null == null
+        assert!(!col.equal_to(0, &input, 2) && !col.equal_to(2, &input, 0));
+
+        let mut buf = true_buf(3);
+        col.vectorized_equal_to(&[0, 1, 2], &input, &[2, 2, 1], &mut buf);
+        assert_eq!(buf_to_vec(&buf), vec![false, false, false]);
+
+        let mut buf = true_buf(3);
+        buf.set_bit(0, false);
+        col.vectorized_equal_to(&[0, 1, 2], &input, &[0, 1, 2], &mut buf);
+        assert_eq!(buf_to_vec(&buf), vec![false, true, true]);
+
+        assert_eq!(col.take_n(0).len(), 0);
+        let taken = col.take_n(3);
+        assert_is_dict_utf8(&taken);
+        assert_eq!(str_values(&taken), vec![None, None, Some("b".into())]);
+
+        let input2 = dict_arr(&[Some(0), None, Some(1), Some(0), None], &["cat", "dog"]);
+        col.vectorized_append(&input2, &[0, 1, 2, 3, 4]).unwrap();
+        let taken2 = col.take_n(5);
+        assert_is_dict_utf8(&taken2);
+        assert_eq!(
+            str_values(&taken2),
+            vec![
+                Some("cat".into()),
+                None,
+                Some("dog".into()),
+                Some("cat".into()),
+                None,
+            ]
+        );
+
+        let empty = Box::new(col).build();
+        assert_is_dict_utf8(&empty);
+        assert_eq!(empty.len(), 0);
+    }
+
+    #[test]
+    fn u64_primitive_inner_multi_batch() {
+        use crate::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder;
+        use arrow::array::UInt64Array;
+        use arrow::datatypes::UInt64Type;
+
+        let field = Field::new("", DataType::UInt64, true);
+        let mut col = DictionaryGroupValuesColumn::<Int32Type>::new(
+            Box::new(PrimitiveGroupValueBuilder::<UInt64Type, true>::new(
+                DataType::UInt64,
+            )),
+            &field,
+        );
+
+        let batch1: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::new(
+            Int32Array::from(vec![Some(0), Some(1), None, Some(2), Some(0)]),
+            Arc::new(UInt64Array::from(vec![10u64, 20, 30])),
+        ));
+        col.vectorized_append(&batch1, &[0, 1, 2, 3, 4]).unwrap();
+
+        let mut buf = true_buf(5);
+        col.vectorized_equal_to(&[0, 1, 2, 3, 4], &batch1, &[0, 1, 2, 3, 0], &mut buf);
+        assert_eq!(buf_to_vec(&buf), vec![true, true, true, true, true]);
+
+        let mut buf = true_buf(2);
+        col.vectorized_equal_to(&[0, 4], &batch1, &[1, 1], &mut buf);
+        assert_eq!(buf_to_vec(&buf), vec![false, false]);
+
+        let u64_val = |arr: &ArrayRef, pos: usize| {
+            let dict = arr.as_dictionary::<Int32Type>();
+            dict.values()
+                .as_any()
+                .downcast_ref::<UInt64Array>()
+                .unwrap()
+                .value(dict.key(pos).unwrap())
+        };
+
+        let taken1 = col.take_n(3);
+        assert_eq!(u64_val(&taken1, 0), 10);
+        assert_eq!(u64_val(&taken1, 1), 20);
+        assert!(taken1.as_dictionary::<Int32Type>().key(2).is_none());
+
+        let batch2: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::new(
+            Int32Array::from(vec![None, Some(0)]),
+            Arc::new(UInt64Array::from(vec![99u64])),
+        ));
+        col.vectorized_append(&batch2, &[0, 1]).unwrap();
+
+        let mut buf = true_buf(2);
+        col.vectorized_equal_to(&[2, 3], &batch2, &[0, 1], &mut buf);
+        assert_eq!(buf_to_vec(&buf), vec![true, true]);
+
+        let out = Box::new(col).build();
+        assert_eq!(u64_val(&out, 0), 30);
+        assert_eq!(u64_val(&out, 1), 10);
+        assert!(out.as_dictionary::<Int32Type>().key(2).is_none());
+        assert_eq!(u64_val(&out, 3), 99);
+    }
+}