jayzhan211 opened a new issue, #26696:
URL: https://github.com/apache/datafusion/issues/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`:
```sql
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):
```rs
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.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
