jayzhan211 opened a new issue, #11885:
URL: https://github.com/apache/arrow-rs/issues/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:
```rust
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);
}
}
```
| values per input dictionary | `concat` of the 12 slices (µs) | the same
rows as `Utf8` (µs) |
|---|---|---|
| 1,000 | 121–141 | 3 |
| 10,000 | 131–134 | 3 |
| 100,000 | 150–153 | 3 |
| 1,000,000 | 426–438 | 3 |
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
- #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 apache/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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
