apache / apache/datafusion

Optimize take_n on DictionaryGroupValuesColumn<K>

Open
#24,089 7 comments 0 reactions 1 assignee Claimed by @saadtajwar View on GitHub
enhancement
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

### Is your feature request related to a problem or challenge?

see context here
- https://github.com/apache/datafusion/pull/23187#discussion_r3688359028

`DictionaryGroupValuesColumn::take_n` emits the first n groups and rebuilds the remainder in-place. Every call hashes all surviving distinct values from scratch to reconstruct value_dedup, making repeated partial emissions O(G² / batch_size) over the lifetime of a query, where G is the total number of distinct values seen. Additionally, `arrow::compute::take` does not compact the backing storage for Utf8View, BinaryView, or nested dictionary value arrays; the rebuilt column retains a reference to the full original allocation rather than releasing memory proportional to the dropped groups. Under a high-cardinality dictionary key with frequent spill-driven partial emissions (e.g. streaming aggregation with a large fan-out), this combination causes both CPU and peak RSS to grow super-linearly with group count.

### Describe the solution you'd like

see context here
- https://github.com/apache/datafusion/pull/23187#discussion_r3708572028

Two targeted improvements:

1. Incremental dedup table on rebuild. Instead of re-hashing all surviving inner slots after a take_n, retain a dirty-flag or generation counter on value_dedup and only remove the entries that were fully emitted (i.e. whose inner_slot is not referenced by any remaining group). This keeps the rebuild cost proportional to the number of emitted distinct values rather than the number of surviving ones, reducing amortized complexity from O(G²/B) to O(G).
2. Storage compaction for view and nested types. After calling `compute::take` to subset the inner values array, explicitly compact Utf8View / BinaryView columns by calling `StringViewArray::gc` (or equivalent) to release unreferenced buffers. For nested dictionary values, apply the same compaction recursively. This bounds peak memory to the live distinct-value set rather than the union of all values ever seen in a partition.

### Describe alternatives you've considered

Keeping the current implementation. This is not ideal due to performance issues mentioned above.

### Additional context

- https://github.com/apache/datafusion/issues/23993
- https://github.com/apache/datafusion/pull/23187
- https://github.com/apache/datafusion/pull/23523

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.