concat of dictionary slices does work proportional to each input's whole dictionary, not its rows #11885

Is your feature request related to a problem or challenge?

When concat merges dictionaries (concat_dictionaries → merge_dictionary_values), it builds state sized to each input's whole values array, even when the input is a short slice that references only a few of those values:

• compute_values_mask(keys, key_mask, values.len()) builds a values.len()-bit mask for each input (arrow-select/src/dictionary.rs:252), and get_masked_values walks it (dictionary.rs:254).
• The key mapping for each input is vec![zero; dictionary.values().len()] (dictionary.rs:272).

So concatenating slices costs time proportional to their dictionaries, not to their rows. This shows up when the inputs are slices of large dictionaries. For example, DataFusion's hash aggregation emits all of its groups as one dictionary and slices it into batch_size batches, and the repartition coalescer then concatenates pieces of those slices from several partitions (apache/datafusion#26696).

The repro below (arrow 60.0.0) concatenates 12 slices of 682 rows each, 8184 rows in total. Each slice comes from a different Dictionary(Int32, Utf8) whose values are all distinct, so the dictionaries are merged:

use arrow_array::types::Int32Type;
use arrow_array::{Array, ArrayRef, DictionaryArray, Int32Array, StringArray};
use arrow_select::concat::concat;
use std::sync::Arc;
use std::time::Instant;

fn dictionary(values: usize, seed: usize) -> DictionaryArray<Int32Type> {
    let v = StringArray::from_iter_values((0..values).map(|i| format!("p{seed}_value_{i}")));
    let keys = Int32Array::from_iter_values(0..values as i32);
    DictionaryArray::try_new(keys, Arc::new(v)).unwrap()
}

fn main() {
    let (pieces, rows) = (12, 682);
    for values in [1_000, 10_000, 100_000, 1_000_000] {
        let dicts: Vec<_> = (0..pieces).map(|p| dictionary(values, p)).collect();
        let slices: Vec<ArrayRef> = dicts
            .iter()
            .map(|d| Arc::new(d.slice((values - rows) / 2, rows)) as ArrayRef)
            .collect();
        let inputs: Vec<&dyn Array> = slices.iter().map(|a| a.as_ref()).collect();
        let start = Instant::now();
        for _ in 0..100 {
            concat(&inputs).unwrap();
        }
        println!("{values}: {:?} per concat", start.elapsed() / 100);
    }
}

The output always has 8184 rows, but the time grows with the size of the input dictionaries. The table shows the best of 15 runs of 20 iterations (release build, Apple M4 Pro, two rounds); the snippet's plain average gives about the same numbers after a warm-up run.

Describe the solution you'd like

When an input has far more values than rows, keep state only for the values its rows reference. For example, collect the referenced keys (O(rows)) and build the mapping and the masked values from them, instead of values.len()-sized buffers. Inputs whose dictionary isn't much larger than their rows keep the current dense path.

Describe alternatives you've considered

• Compacting oversized inputs before merging (take(values, keys) with keys 0..len) in concat_dictionaries. I tried this: it brings the 1,000,000 row down to 155 µs, but it makes the small cases 15–25% slower (150–164 µs at 1,000 values). Copying the values costs more than the dense mapping until the dictionary is very large, so this would need a high threshold.
• garbage_collect_dictionary on each input first. It also allocates a values.len() remap per input (dictionary.rs:59).
• Compacting in each caller before concatenating. This works, but every caller that slices dictionaries would have to do it.

Additional context

• perf: gc dictionaries in interleave_fallback_dictionary #9939 is the same kind of problem in interleave_fallback_dictionary, the path that doesn't merge. interleave_dictionaries (interleave.rs:304) also calls merge_dictionary_values and keeps the same per-input values.len() state.
• In the DataFusion query from Hash aggregate output slices keep the whole dictionary, so downstream work grows with the group count datafusion#26696, each piece carries about 600k values (7.25M partial groups over 12 partitions). Compacting oversized inputs in concat_dictionaries took the repartition's send_time from 3.73 s to 1.94 s (summed over 12 partitions) and the query from 878 ms to 708 ms.
• Merging is slower than concatenating the equivalent Utf8 even for small dictionaries (about 130 µs vs 3 µs above), because every referenced value is interned. This issue is only about the part that grows with the dictionary size.

I'm not planning to fix this myself, so anyone is welcome to pick it up. Feel free to ping me with questions about the cause or when you'd like a review.