Hash aggregate output slices keep the whole dictionary, so downstream work grows with the group count #26696

Is your feature request related to a problem or challenge?

When a hash aggregate groups on a dictionary column, it materializes all groups once and then emits them in batch_size slices (MaterializedAggregateOutput::next_batch in aggregates/aggregate_hash_table/common.rs). Slicing a DictionaryArray keeps its whole values array, so every 8192-row slice still carries the values of every group in that partition. The operators after it then do work proportional to the group count instead of the slice:

• the hash RepartitionExec hashes each slice with hash_dictionary, which hashes every dictionary value
• its output coalescer concatenates the pieces with merge_dictionary_values, which allocates and scans per-value state for each piece
• the final aggregate hashes the dictionary values again

On ClickBench hits_partitioned with dictionary reads enabled, grouping on a high-cardinality dictionary column is 2.6x slower than grouping on the same column cast to Utf8:

SET datafusion.execution.parquet.binary_as_string = true;
SET datafusion.execution.parquet.enable_rle_to_dictionary = true;
CREATE EXTERNAL TABLE hits STORED AS PARQUET LOCATION 'benchmarks/data/hits_partitioned/';

-- 1339 ms
SELECT "SearchPhrase" s, COUNT(*) c FROM hits WHERE "SearchPhrase" <> '' GROUP BY s ORDER BY c DESC LIMIT 10;
-- 520 ms
SELECT arrow_cast("SearchPhrase", 'Utf8') s, COUNT(*) c FROM hits WHERE "SearchPhrase" <> '' GROUP BY s ORDER BY c DESC LIMIT 10;

In EXPLAIN ANALYZE of the first query, the hash RepartitionExec reports repartition_time=3.78s and send_time=3.14s (summed over 12 partitions).

Describe the solution you'd like

When an emitted slice's dictionary has more values than the slice has rows, give the slice its own values. Group keys are distinct, so take(values, keys) plus sequential keys is enough, and it costs O(rows in the slice):

fn compact_dictionary<K: ArrowDictionaryKeyType>(
    dictionary: &DictionaryArray<K>,
) -> Result<ArrayRef> {
    let values = take(dictionary.values().as_ref(), dictionary.keys(), None)?;
    let keys = (0..dictionary.len())
        .map(|i| {
            K::Native::from_usize(i).ok_or_else(|| {
                internal_datafusion_err!("dictionary slice longer than its key type")
            })
        })
        .collect::<Result<Vec<_>>>()?;
    let keys = PrimitiveArray::<K>::new(keys.into(), dictionary.keys().nulls().cloned());
    Ok(Arc::new(DictionaryArray::<K>::try_new(keys, values)?))
}

Applying this to each slice returned by next_batch brings the first query from 1339 ms to 754 ms (min of 3 rounds, 12 partitions). Queries whose group keys aren't dictionaries are unaffected, since the check is a no-op for them. The aggregate unit tests and the aggregate*, group_by*, dictionary*, distinct* and repartition* sqllogictest files pass with it.

Describe alternatives you've considered

• Hashing only the referenced values in hash_dictionary when a batch has fewer rows than values. That fixes the repartition hashing, but not the coalescer merge or the final aggregate.
• arrow::compute::garbage_collect_any_dictionary on each slice. It's shorter, but it allocates a remap table the size of the whole dictionary per slice, so the cost still grows with the group count.

Additional context

Even with this change, the first query stays slower than the Utf8 one (754 ms vs 520 ms). #24822 tracks the dictionary group-key cost more generally.

This came up while reviewing #25874, where substr on a dictionary column now returns a dictionary, so grouping on its result goes through this path. There, this change fixes the 2-partition case (1.30x slower than main becomes 1.03x), but with 12 partitions the query stays about 1.2x slower. The remaining time is in the repartition coalescer, which still rebuilds a dictionary for each output batch (send_time 279 ms vs 41 ms with Utf8 keys), and in the final aggregate.

Timings are local, from a release-nonlto build, best of 3 rounds.