Skip to content

Optimize take_n on DictionaryGroupValuesColumn<K> #24089

Description

@Rich-T-kid

Is your feature request related to a problem or challenge?

see context here

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

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

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions