add regression test to align with GroupValueRows & re-order overflow check

Author: Rich-T-kid

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

@@ -58,12 +58,12 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> DictionaryGroupValuesColumn<K> {
 
     /// 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) {
+    fn check_key_overflow(current_len: usize) -> Result<()> {
+        if !Self::valid_bounds::<K>(current_len) {
             return exec_err!(
                 "Dictionary key type {:?} cannot represent {} groups",
                 K::DATA_TYPE,
-                current_len + additional
+                current_len
             );
         }
         Ok(())
@@ -97,12 +97,12 @@ 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),
-            Some(key) => self.inner.append_val(dict.values(), key),
+            None => self.inner.append_val(&self.null_array, 0)?,
+            Some(key) => self.inner.append_val(dict.values(), key)?,
         }
+        Self::check_key_overflow(self.inner.len())
     }
 
     fn vectorized_equal_to(
@@ -161,7 +161,6 @@ 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();
 
@@ -170,15 +169,17 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
                 .iter()
                 .map(|&row_index| dict_keys.value(row_index).as_usize())
                 .collect();
-            self.inner.vectorized_append(dict.values(), &value_indices)
+            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)
+            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 {
@@ -187,8 +188,8 @@ impl<K: ArrowDictionaryKeyType + Send + Sync> GroupColumn
                     Some(key) => self.inner.append_val(dict.values(), key)?,
                 }
             }
-            Ok(())
         }
+        Self::check_key_overflow(self.inner.len())
     }
 
     fn len(&self) -> usize {
@@ -370,4 +371,37 @@ mod tests {
         assert!(out.as_dictionary::<Int32Type>().key(2).is_none());
         assert_eq!(u64_val(&out, 3), 99);
     }
+
+    // Regression test for https://github.com/apache/datafusion/issues/23127 —
+    // we should fail early when we can't produce the correct output schema.
+    #[test]
+    fn uint8_key_overflow() {
+        use arrow::array::UInt8Array;
+        use arrow::datatypes::UInt8Type;
+        use datafusion_physical_expr::binary_map::OutputType;
+
+        let field = Field::new("", DataType::Utf8, true);
+        let mut col = DictionaryGroupValuesColumn::<UInt8Type>::new(
+            Box::new(ByteGroupValueBuilder::<i32>::new(OutputType::Utf8)),
+            &field,
+        );
+
+        // fill all 256 distinct values (u8 key range 0..=255)
+        let strs: Vec<String> = (0..=255u16).map(|i| i.to_string()).collect();
+        let str_refs: Vec<&str> = strs.iter().map(|s| s.as_str()).collect();
+        let full: ArrayRef = Arc::new(DictionaryArray::<UInt8Type>::new(
+            UInt8Array::from((0..=255u8).map(Some).collect::<Vec<_>>()),
+            Arc::new(StringArray::from(str_refs)),
+        ));
+        let rows: Vec<usize> = (0..256).collect();
+        col.vectorized_append(&full, &rows).unwrap();
+        assert_eq!(col.len(), 256);
+
+        // one more must overflow the key range
+        let extra: ArrayRef = Arc::new(DictionaryArray::<UInt8Type>::new(
+            UInt8Array::from(vec![Some(0u8)]),
+            Arc::new(StringArray::from(vec!["overflow"])),
+        ));
+        assert!(col.append_val(&extra, 0).is_err());
+    }
 }