perf: bound work over dictionary values by the batch's row count

clippy.toml

@@ -1,6 +1,8 @@
 disallowed-methods = [
     { path = "tokio::task::spawn", reason = "To provide cancel-safety, use `SpawnedTask::spawn` instead (https://github.com/apache/datafusion/issues/6513)" },
     { path = "tokio::task::spawn_blocking", reason = "To provide cancel-safety, use `SpawnedTask::spawn_blocking` instead (https://github.com/apache/datafusion/issues/6513)" },
+    { path = "arrow_array::array::AnyDictionaryArray::with_values", reason = "Running code on every dictionary value can cost far more than the batch has rows; use `datafusion_common::utils::map_dictionary_values`" },
+    { path = "arrow_array::array::DictionaryArray::with_values", reason = "Running code on every dictionary value can cost far more than the batch has rows; use `datafusion_common::utils::map_dictionary_values`" },
     { path = "std::vec::Vec::reserve", reason = "Use `Vec::try_reserve` so allocation failures can be reported instead of panicking", replacement = "try_reserve" },
 ]
 

datafusion/common/Cargo.toml

@@ -72,6 +72,7 @@ name = "record_batch_memory"
 arrow = { workspace = true }
 arrow-ipc = { workspace = true }
 arrow-schema = { workspace = true, features = ["canonical_extension_types"] }
+arrow-select = { workspace = true }
 chrono = { workspace = true }
 foldhash = "0.2"
 half = { workspace = true }

datafusion/common/benches/with_hashes.rs

@@ -20,7 +20,7 @@
 use arrow::array::{
     Array, ArrayRef, ArrowPrimitiveType, DictionaryArray, FixedSizeListArray,
     GenericStringArray, Int32Array, Int64Array, ListArray, ListViewArray, MapArray,
-    NullArray, NullBufferBuilder, OffsetSizeTrait, PrimitiveArray, RunArray,
+    NullArray, NullBufferBuilder, OffsetSizeTrait, PrimitiveArray, RunArray, StringArray,
     StringViewArray, StructArray, UnionArray, make_array,
 };
 use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
@@ -407,6 +407,22 @@ fn sliced_array_benchmark(c: &mut Criterion) {
             );
         }
     }
+
+    // Sliced DictionaryArray with one value per row of the full array, like the
+    // group keys an aggregation emits in `batch_size` slices. Hashing reduces
+    // it to the referenced values only above 16 values per row, so 1/8 keeps
+    // the full dictionary and 1/32 doesn't.
+    for ratio in [8, 32] {
+        let total_rows = BATCH_SIZE * ratio;
+        let full_array = distinct_dictionary_array(total_rows);
+        let sliced = full_array.slice(BATCH_SIZE * (ratio / 2), BATCH_SIZE);
+        c.bench_function(
+            &format!("dictionary_sliced: 1/{ratio} of {total_rows} values"),
+            |b| {
+                do_hash_test_with_len(b, std::slice::from_ref(&sliced), BATCH_SIZE);
+            },
+        );
+    }
 }
 
 fn do_hash_test_with_len(b: &mut Bencher, arrays: &[ArrayRef], expected_len: usize) {
@@ -420,6 +436,14 @@ fn do_hash_test_with_len(b: &mut Bencher, arrays: &[ArrayRef], expected_len: usi
     });
 }
 
+/// A `Dictionary(Int32, Utf8)` with a distinct value for every row
+fn distinct_dictionary_array(num_rows: usize) -> ArrayRef {
+    let values =
+        StringArray::from_iter_values((0..num_rows).map(|i| format!("group_{i}")));
+    let keys = Int32Array::from_iter_values(0..num_rows as i32);
+    Arc::new(DictionaryArray::<Int32Type>::try_new(keys, Arc::new(values)).unwrap())
+}
+
 fn list_array(num_rows: usize) -> ArrayRef {
     let mut rng = make_rng();
     let elements_per_row = 5;

datafusion/common/src/hash_utils.rs

@@ -17,6 +17,8 @@
 
 //! Functionality used both on logical and physical plans
 
+#[cfg(not(feature = "force_hash_collisions"))]
+use crate::utils::compact_oversized_dictionary;
 use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano};
 use arrow::array::*;
 #[cfg(not(feature = "force_hash_collisions"))]
@@ -671,7 +673,11 @@ fn hash_dictionary<K: ArrowDictionaryKeyType>(
 ) -> Result<()> {
     // Hash each dictionary value once, and then use that computed
     // hash for each key value to avoid a potentially expensive
-    // redundant hashing for large dictionary elements (e.g. strings)
+    // redundant hashing for large dictionary elements (e.g. strings).
+    // A dictionary with many more values than rows (e.g. a slice) is first
+    // reduced to the values its rows reference.
+    let compacted = compact_oversized_dictionary(array)?;
+    let array = compacted.as_ref().unwrap_or(array);
     let dict_values = array.values();
     let mut dict_hashes = vec![0; dict_values.len()];
     create_hashes([dict_values], random_state, &mut dict_hashes)?;
@@ -686,6 +692,8 @@ fn hash_dictionary_with_child_hashing<K: ArrowDictionaryKeyType>(
     hashes_buffer: &mut [u64],
     multi_col: bool,
 ) -> Result<()> {
+    let compacted = compact_oversized_dictionary(array)?;
+    let array = compacted.as_ref().unwrap_or(array);
     let dict_values = array.values();
     let mut dict_hashes = vec![0; dict_values.len()];
     child_hashing.create_hashes([dict_values], &mut dict_hashes)?;
@@ -1851,6 +1859,42 @@ mod tests {
         assert_ne!(dict_hashes[0], dict_hashes[2]);
     }
 
+    #[test]
+    // Tests actual values of hashes, which are different if forcing collisions
+    #[cfg(not(feature = "force_hash_collisions"))]
+    fn create_hashes_for_dict_slices_with_more_values_than_rows() {
+        use crate::utils::MAX_DICTIONARY_VALUES_PER_ROW;
+
+        let values: ArrayRef = Arc::new(StringArray::from_iter(
+            (0..40).map(|i| (i != 2).then(|| format!("v{i}"))),
+        ));
+        let keys = Int8Array::from_iter((0..48).map(|i| (i % 7 != 1).then_some(i % 40)));
+        let dict: ArrayRef = Arc::new(DictionaryArray::try_new(keys, values).unwrap());
+        let ints: ArrayRef = Arc::new(Int32Array::from_iter_values(0..48));
+        let random_state = RandomState::with_seed(0);
+
+        let mut single = vec![0; dict.len()];
+        create_hashes([&dict], &random_state, &mut single).unwrap();
+        let mut multi = vec![0; dict.len()];
+        create_hashes([&ints, &dict], &random_state, &mut multi).unwrap();
+
+        for (offset, len) in [(0, 2), (1, 2), (8, 1), (46, 2)] {
+            let slice = dict.slice(offset, len);
+            assert!(
+                slice.as_any_dictionary().values().len()
+                    > MAX_DICTIONARY_VALUES_PER_ROW * len
+            );
+            let mut hashes = vec![0; len];
+            create_hashes([&slice], &random_state, &mut hashes).unwrap();
+            assert_eq!(hashes, single[offset..offset + len]);
+
+            let ints = ints.slice(offset, len);
+            let mut hashes = vec![0; len];
+            create_hashes([&ints, &slice], &random_state, &mut hashes).unwrap();
+            assert_eq!(hashes, multi[offset..offset + len]);
+        }
+    }
+
     #[test]
     // Tests actual values of hashes, which are different if forcing collisions
     #[cfg(not(feature = "force_hash_collisions"))]

datafusion/common/src/utils/mod.rs

@@ -26,25 +26,29 @@ pub mod string_utils;
 
 use crate::assert_or_internal_err;
 use crate::error::{
-    _exec_datafusion_err, _exec_err, _internal_datafusion_err, _plan_datafusion_err,
+    _exec_datafusion_err, _exec_err, _internal_datafusion_err, _internal_err,
+    _plan_datafusion_err,
 };
 use crate::{Result, ScalarValue};
 use arrow::array::{
     Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, OffsetSizeTrait,
     cast::AsArray,
 };
 use arrow::array::{
-    ArrowPrimitiveType, BooleanArray, Datum, GenericListArray, Int32Array, Int64Array,
-    MutableArrayData, PrimitiveArray, layout, make_array,
+    ArrowPrimitiveType, BooleanArray, Datum, DictionaryArray, GenericListArray,
+    Int32Array, Int64Array, MutableArrayData, PrimitiveArray, layout, make_array,
 };
 use arrow::array::{LargeListViewArray, ListViewArray};
 use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
 use arrow::compute::kernels::cmp::eq;
 use arrow::compute::kernels::length::length;
-use arrow::compute::{SortColumn, SortOptions, nullif, partition};
+use arrow::compute::{SortColumn, SortOptions, nullif, partition, take};
 use arrow::datatypes::{
-    ArrowNativeType, DataType, Field, Int32Type, Int64Type, Metadata, SchemaRef,
+    ArrowDictionaryKeyType, ArrowNativeType, DataType, Field, Int32Type, Int64Type,
+    Metadata, SchemaRef,
 };
+use arrow::downcast_dictionary_array;
+use arrow_select::dictionary::garbage_collect_any_dictionary;
 #[cfg(feature = "sql")]
 use sqlparser::{ast::Ident, dialect::GenericDialect, parser::Parser};
 use std::borrow::{Borrow, Cow};
@@ -416,6 +420,70 @@ pub fn split_vec_min_alloc<T>(vec: &mut Vec<T>, n: usize) -> Vec<T> {
     }
 }
 
+/// Applies `f` to the values of the dictionary `array` and keeps its keys.
+///
+/// Slices, filters and repartitioned batches keep the whole values array of
+/// the dictionary they came from, so running `f` on every value can cost far
+/// more than the batch has rows. `f` never runs on more than
+/// [`MAX_DICTIONARY_VALUES_PER_ROW`] values per row: past that, it runs on the
+/// values the rows reference instead.
+///
+/// `f` can also see values that no row references. If it fails, it runs once
+/// more on only the referenced values, so it fails only when running it on
+/// each row would.
+pub fn map_dictionary_values<F>(array: &ArrayRef, mut f: F) -> Result<ArrayRef>
+where
+    F: FnMut(&ArrayRef) -> Result<ArrayRef>,
+{
+    let compacted = downcast_dictionary_array! {
+        array => compact_oversized_dictionary(array)?.map(|d| Arc::new(d) as ArrayRef),
+        dt => return _internal_err!("expected a dictionary array, got {dt}"),
+    };
+    let mut dictionary = compacted.unwrap_or_else(|| Arc::clone(array));
+    let values = match f(dictionary.as_any_dictionary().values()) {
+        Ok(values) => values,
+        Err(e) => {
+            let referenced =
+                garbage_collect_any_dictionary(dictionary.as_any_dictionary())?;
+            if referenced.as_any_dictionary().values().len()
+                == dictionary.as_any_dictionary().values().len()
+            {
+                return Err(e);
+            }
+            dictionary = referenced;
+            f(dictionary.as_any_dictionary().values())?
+        }
+    };
+    #[expect(clippy::disallowed_methods)] // the values are bounded above
+    Ok(dictionary.as_any_dictionary().with_values(values))
+}
+
+/// The most dictionary values per row that code processing every value of a
+/// dictionary should handle. Past this, taking the values the rows reference
+/// is cheaper than processing all of them, even for the cheapest per-value
+/// work: the break-even is about 4 values per row for hashing and about 13 for
+/// ASCII `lower`/`upper`, which convert the whole values buffer at once. Any
+/// bound removes the cost that grows with the dictionary size; the exact value
+/// only trades a constant factor.
+pub const MAX_DICTIONARY_VALUES_PER_ROW: usize = 16;
+
+/// Returns `None` when `array` has at most [`MAX_DICTIONARY_VALUES_PER_ROW`]
+/// values per row, and otherwise a copy whose values are `take(values, keys)`
+/// and whose keys are `0..len`, so that work over its values costs O(rows).
+pub(crate) fn compact_oversized_dictionary<K: ArrowDictionaryKeyType>(
+    array: &DictionaryArray<K>,
+) -> Result<Option<DictionaryArray<K>>> {
+    if array.values().len() <= MAX_DICTIONARY_VALUES_PER_ROW * array.len() {
+        return Ok(None);
+    }
+    let values = take(array.values().as_ref(), array.keys(), None)?;
+    // `len < values.len()`, so every row index fits in the key type
+    let keys: ScalarBuffer<K::Native> =
+        (0..array.len()).map(K::Native::usize_as).collect();
+    let keys = PrimitiveArray::<K>::new(keys, array.keys().nulls().cloned());
+    Ok(Some(DictionaryArray::try_new(keys, values)?))
+}
+
 #[cfg(test)]
 mod split_vec_min_alloc_tests {
     use super::split_vec_min_alloc;
@@ -1501,14 +1569,15 @@ pub fn apply_parent_nulls(
 
 #[cfg(test)]
 mod tests {
+    use std::cell::Cell;
     use std::collections::HashMap;
     use std::sync::Arc;
 
     use super::*;
     use crate::ScalarValue::Null;
     use arrow::{
-        array::{Float64Array, Int32Array, NullArray},
-        datatypes::Int32Type,
+        array::{Float64Array, Int32Array, NullArray, StringArray},
+        datatypes::{Int8Type, Int32Type},
     };
     #[cfg(feature = "sql")]
     use sqlparser::ast::Ident;
@@ -2106,4 +2175,71 @@ mod tests {
         fsl_values_row_number(-1, 2).unwrap_err();
         fsl_values_row_number(-1, 0).unwrap_err();
     }
+
+    #[test]
+    fn map_dictionary_values_bounds_values_per_row() -> Result<()> {
+        let strings = (0..24)
+            .map(|i| (i % 5 != 1).then(|| "x".repeat(i)))
+            .collect::<Vec<_>>();
+        let dict: DictionaryArray<Int32Type> =
+            strings.iter().map(|s| s.as_deref()).collect();
+        let dict: ArrayRef = Arc::new(dict);
+        for (offset, len) in [(0, 24), (1, 1), (6, 1)] {
+            let slice = dict.slice(offset, len);
+            let mut seen = 0;
+            let mapped = map_dictionary_values(&slice, |values| {
+                seen = values.len();
+                Ok(length(values)?)
+            })?;
+            assert!(
+                seen <= MAX_DICTIONARY_VALUES_PER_ROW * len,
+                "{seen} values for {len} rows"
+            );
+            let expected = length(&arrow::compute::cast(&slice, &DataType::Utf8)?)?;
+            assert_eq!(&arrow::compute::cast(&mapped, &DataType::Int32)?, &expected);
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn map_dictionary_values_fails_only_on_referenced_values() -> Result<()> {
+        let calls = Cell::new(0);
+        let fail_on_bad = |values: &ArrayRef| {
+            calls.set(calls.get() + 1);
+            if values.as_string::<i32>().iter().any(|v| v == Some("bad")) {
+                return _exec_err!("bad value");
+            }
+            Ok(length(values)?)
+        };
+        let values: ArrayRef =
+            Arc::new(StringArray::from(vec!["ok", "bad", "fine", "bad"]));
+
+        // "bad" at 1 is behind null keys only, and at 3 has no key. There are
+        // more rows than Int8 keys can index, so the retry can't give each row
+        // its own key.
+        let keys = (0..200)
+            .map(|i| [1, 0, 2][i % 3])
+            .collect::<ScalarBuffer<i8>>();
+        let nulls = NullBuffer::from_iter((0..200).map(|i| i % 3 != 0));
+        let keys = PrimitiveArray::<Int8Type>::new(keys, Some(nulls));
+        let dict: ArrayRef =
+            Arc::new(DictionaryArray::try_new(keys, Arc::clone(&values))?);
+        let mapped = map_dictionary_values(&dict, fail_on_bad)?;
+        assert_eq!(
+            mapped.data_type(),
+            &DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32))
+        );
+        let expected = length(&arrow::compute::cast(&dict, &DataType::Utf8)?)?;
+        assert_eq!(&arrow::compute::cast(&mapped, &DataType::Int32)?, &expected);
+        assert_eq!(calls.get(), 2);
+
+        // A row references "bad", so `f` must still fail, without a retry when
+        // every value is referenced
+        calls.set(0);
+        let keys = PrimitiveArray::<Int8Type>::from(vec![0, 1, 2, 3]);
+        let dict: ArrayRef = Arc::new(DictionaryArray::try_new(keys, values)?);
+        map_dictionary_values(&dict, fail_on_bad).unwrap_err();
+        assert_eq!(calls.get(), 1);
+        Ok(())
+    }
 }

datafusion/functions-aggregate/src/array_agg.rs

@@ -36,6 +36,7 @@ use arrow_select::dictionary::garbage_collect_any_dictionary;
 use datafusion_common::cast::as_list_array;
 use datafusion_common::hash_utils::{RandomState, create_hashes};
 use datafusion_common::scalar::copy_array_data;
+use datafusion_common::utils::map_dictionary_values;
 use datafusion_common::utils::memory::estimate_memory_size;
 use datafusion_common::utils::proxy::HashTableAllocExt;
 use datafusion_common::utils::{SingleRowListArrayBuilder, offset_span};
@@ -2061,11 +2062,11 @@ fn compact_payload(array: ArrayRef) -> Result<ArrayRef> {
         DataType::BinaryView => Ok(Arc::new(array.as_binary_view().gc())),
         DataType::Dictionary(_, _) => {
             let array = garbage_collect_any_dictionary(array.as_any_dictionary())?;
-            let dictionary = array.as_any_dictionary();
             // GC can return the original values if every entry is referenced.
             // Detach offset slices as well as any nested view/dictionary buffers.
-            let values = make_array(copy_array_data(&dictionary.values().to_data()));
-            Ok(dictionary.with_values(compact_payload(values)?))
+            map_dictionary_values(&array, |values| {
+                compact_payload(make_array(copy_array_data(&values.to_data())))
+            })
         }
         data_type if data_type.is_nested() => {
             let data = array.to_data();

datafusion/functions/src/core/getfield.rs

@@ -121,6 +121,8 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result<Column
                     )
                 })?;
             let field_col = apply_parent_nulls(field_col, values_struct.nulls())?;
+            #[expect(clippy::disallowed_methods)]
+            // picks a child column, no per-value work
             Ok(ColumnarValue::Array(dict.with_values(field_col)))
         }
         (DataType::Map(_, _), key, _) => {