Feat:[Draft][22715] add dictionaries as a supported group column type

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

@@ -0,0 +1,386 @@
+// 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::compute::concat;
+use arrow::datatypes::DataType;
+use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, Field};
+use datafusion_common::Result;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+pub struct DictionaryGroupValuesColumn<K: ArrowDictionaryKeyType + Send + Sync> {
+    inner: Box<dyn GroupColumn>,
+    null_array: ArrayRef,
+    cached_values: Option<ArrayRef>,
+    cached_combined: Option<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,
+            cached_values: None,
+            cached_combined: None,
+            _phantom: PhantomData,
+        }
+    }
+
+    #[inline]
+    fn get_combined(&mut self, values: &ArrayRef) -> Result<ArrayRef> {
+        let is_cached = self
+            .cached_values
+            .as_ref()
+            .is_some_and(|cached| Arc::ptr_eq(cached, values));
+        if !is_cached {
+            self.cached_combined =
+                Some(concat(&[values.as_ref(), self.null_array.as_ref()])?);
+            self.cached_values = Some(Arc::clone(values));
+        }
+        Ok(Arc::clone(self.cached_combined.as_ref().unwrap()))
+    }
+
+    #[inline]
+    fn into_dict(values: ArrayRef) -> ArrayRef {
+        // at some point in the future we may want to support bumping the key types
+        // https://github.com/apache/datafusion/issues/23127
+        let num_values = values.len();
+        assert!(
+            Self::valid_bounds::<K>(num_values),
+            "Dictionary key type {:?} cannot hold {} values",
+            K::DATA_TYPE,
+            num_values
+        );
+        let keys: PrimitiveArray<K> = (0..num_values)
+            .map(|idx| {
+                if values.is_null(idx) {
+                    None
+                } else {
+                    Some(K::Native::usize_as(idx))
+                }
+            })
+            .collect();
+        Arc::new(DictionaryArray::<K>::new(keys, values))
+    }
+    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<()> {
+        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 keys = dict.keys();
+        let values = dict.values();
+
+        if keys.null_count() == 0 {
+            let key_indices: Vec<usize> = rhs_rows
+                .iter()
+                .map(|&row| keys.value(row).as_usize())
+                .collect();
+            self.inner.vectorized_equal_to(
+                lhs_rows,
+                values,
+                &key_indices,
+                equal_to_results,
+            );
+        } else {
+            for (idx, (lhs_row, rhs_row)) in lhs_rows.iter().zip(rhs_rows).enumerate() {
+                if !equal_to_results.get_bit(idx) {
+                    continue;
+                }
+                let is_equal = match dict.key(*rhs_row) {
+                    None => self.inner.equal_to(*lhs_row, &self.null_array, 0),
+                    Some(key) => self.inner.equal_to(*lhs_row, values, key),
+                };
+                if !is_equal {
+                    equal_to_results.set_bit(idx, false);
+                }
+            }
+        }
+    }
+
+    fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> {
+        let dict = array.as_dictionary::<K>();
+        let keys = dict.keys();
+
+        if keys.null_count() == 0 {
+            let key_indices: Vec<usize> =
+                rows.iter().map(|&row| keys.value(row).as_usize()).collect();
+            self.inner.vectorized_append(dict.values(), &key_indices)
+        } else {
+            let combined = self.get_combined(dict.values())?;
+            // last element of combined is always the null sentinel appended by get_combined
+            let null_idx = combined.len() - 1;
+            let key_indices: Vec<usize> = rows
+                .iter()
+                .map(|&row| dict.key(row).unwrap_or(null_idx))
+                .collect();
+            self.inner.vectorized_append(&combined, &key_indices)
+        }
+    }
+
+    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 string_arr = plain.as_any().downcast_ref::<StringArray>().unwrap();
+        (0..string_arr.len())
+            .map(|idx| {
+                string_arr
+                    .is_valid(idx)
+                    .then(|| string_arr.value(idx).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 builder = BooleanBufferBuilder::new(len);
+        builder.append_n(len, true);
+        builder
+    }
+
+    fn buf_to_vec(buf: &BooleanBufferBuilder) -> Vec<bool> {
+        (0..buf.len()).map(|idx| buf.get_bit(idx)).collect()
+    }
+
+    #[test]
+    fn utf8_dict_null_handling() {
+        let mut col = utf8_col();
+        let null_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(&null_input, row).unwrap();
+        }
+
+        assert!(
+            col.equal_to(0, &null_input, 0)
+                && col.equal_to(1, &null_input, 1)
+                && col.equal_to(0, &null_input, 1)
+        );
+        assert!(!col.equal_to(0, &null_input, 2) && !col.equal_to(2, &null_input, 0));
+
+        let mut eq_results = true_buf(3);
+        col.vectorized_equal_to(&[0, 1, 2], &null_input, &[2, 2, 1], &mut eq_results);
+        assert_eq!(buf_to_vec(&eq_results), vec![false, false, false]);
+
+        let mut preseed_results = true_buf(3);
+        preseed_results.set_bit(0, false);
+        col.vectorized_equal_to(
+            &[0, 1, 2],
+            &null_input,
+            &[0, 1, 2],
+            &mut preseed_results,
+        );
+        assert_eq!(buf_to_vec(&preseed_results), vec![false, true, true]);
+
+        assert_eq!(col.take_n(0).len(), 0);
+        let first_taken = col.take_n(3);
+        assert_is_dict_utf8(&first_taken);
+        assert_eq!(str_values(&first_taken), vec![None, None, Some("b".into())]);
+
+        let second_input =
+            dict_arr(&[Some(0), None, Some(1), Some(0), None], &["cat", "dog"]);
+        col.vectorized_append(&second_input, &[0, 1, 2, 3, 4])
+            .unwrap();
+        let second_taken = col.take_n(5);
+        assert_is_dict_utf8(&second_taken);
+        assert_eq!(
+            str_values(&second_taken),
+            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::{DataType, UInt64Type};
+
+        let field = Field::new("", DataType::UInt64, true);
+        let mut col = DictionaryGroupValuesColumn::<Int32Type>::new(
+            Box::new(PrimitiveGroupValueBuilder::<UInt64Type, true>::new(
+                DataType::UInt64,
+            )),
+            &field,
+        );
+
+        let first_batch: 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(&first_batch, &[0, 1, 2, 3, 4])
+            .unwrap();
+
+        let mut eq_results = true_buf(5);
+        col.vectorized_equal_to(
+            &[0, 1, 2, 3, 4],
+            &first_batch,
+            &[0, 1, 2, 3, 0],
+            &mut eq_results,
+        );
+        assert_eq!(buf_to_vec(&eq_results), vec![true, true, true, true, true]);
+
+        let mut mismatch_results = true_buf(2);
+        col.vectorized_equal_to(&[0, 4], &first_batch, &[1, 1], &mut mismatch_results);
+        assert_eq!(buf_to_vec(&mismatch_results), vec![false, false]);
+
+        let first_taken = col.take_n(3);
+        let first_dict = first_taken.as_dictionary::<Int32Type>();
+        let first_values = first_dict
+            .values()
+            .as_any()
+            .downcast_ref::<UInt64Array>()
+            .unwrap();
+        assert_eq!(first_values.value(first_dict.key(0).unwrap()), 10);
+        assert_eq!(first_values.value(first_dict.key(1).unwrap()), 20);
+        assert!(first_dict.key(2).is_none());
+
+        let second_batch: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::new(
+            Int32Array::from(vec![None, Some(0)]),
+            Arc::new(UInt64Array::from(vec![99u64])),
+        ));
+        col.vectorized_append(&second_batch, &[0, 1]).unwrap();
+
+        let mut cross_batch_results = true_buf(2);
+        col.vectorized_equal_to(
+            &[2, 3],
+            &second_batch,
+            &[0, 1],
+            &mut cross_batch_results,
+        );
+        assert_eq!(buf_to_vec(&cross_batch_results), vec![true, true]);
+
+        let final_output = Box::new(col).build();
+        let final_dict = final_output.as_dictionary::<Int32Type>();
+        let final_values = final_dict
+            .values()
+            .as_any()
+            .downcast_ref::<UInt64Array>()
+            .unwrap();
+        assert_eq!(final_values.value(final_dict.key(0).unwrap()), 30);
+        assert_eq!(final_values.value(final_dict.key(1).unwrap()), 10);
+        assert!(final_dict.key(2).is_none());
+        assert_eq!(final_values.value(final_dict.key(3).unwrap()), 99);
+    }
+}