revised PR comments

Author: Rich-T-kid

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

@@ -20,7 +20,7 @@ use arrow::array::{
     Array, ArrayRef, AsArray, BooleanBufferBuilder, DictionaryArray, PrimitiveArray,
 };
 use arrow::datatypes::{ArrowDictionaryKeyType, ArrowNativeType, DataType, Field};
-use datafusion_common::Result;
+use datafusion_common::{Result, exec_err};
 use std::marker::PhantomData;
 use std::sync::Arc;
 
@@ -45,12 +45,6 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> {
 
     fn into_dict(values: ArrayRef) -> ArrayRef {
         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(|i| (!values.is_null(i)).then(|| K::Native::usize_as(i)))
             .collect();
@@ -59,8 +53,20 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> {
 
     /// 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.null_count() > 0)
-            .then(|| (0..values.len()).find(|&i| values.is_null(i)).unwrap())
+        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 {
@@ -91,6 +97,7 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
     }
 
     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),
@@ -154,6 +161,7 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
     }
 
     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();