apache / apache/datafusion

Add arbitrary group retention for grouped aggregation state

Open
#24,603 3 comments 0 reactions 1 assignee Claimed by @rkrishn7 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?

Long-lived aggregation state needs to delete groups that are no longer active. A streaming engine may evict keys after a TTL or watermark while keeping every other group available for future updates.

`EmitTo::First(n)` can remove only a prefix, and it emits values as part of the operation. Arbitrary cleanup currently requires draining all `G` group keys and accumulator states, selecting the `K` survivors, rebuilding the group interner, and merging every survivor back into each accumulator. This is expensive and temporarily materializes the complete state.

For example, expiring groups `{1, 4, 7}` from ten live groups cannot be expressed without rebuilding all remaining groups.

### Describe the solution you'd like

Add `retain_groups` to both state-owning traits:

```rust
pub trait GroupValues: Send {
fn retain_groups(&mut self, retained_group_indices: &[usize]) -> Result<()>;
}

pub trait GroupsAccumulator: Send + Any {
fn retain_groups(
&mut self,
retained_group_indices: &[usize],
total_num_groups: usize,
) -> Result<()>;
}
```

The input order defines the new dense group IDs. Retaining `[5, 2]` maps old group `5` to new group `0` and old group `2` to new group `1`. Indices must be unique and in bounds. An empty selection clears the state; retaining every group in order is a no-op.

The operation destroys the complement and returns no output. A caller applies the same selection to `GroupValues` and every accumulator before using the new IDs. This keeps retention separate from `EmitTo`, whose contract combines output with prefix removal.

Universal defaults make the API usable by custom implementations:

- `GroupsAccumulator` can call `state(EmitTo::All)`, use Arrow `take`, and merge the retained state under new group IDs.
- `GroupValues` can materialize the selected keys using the preserving-read API, clear its lookup state, and re-intern the survivors. An emit-and-take fallback can support implementations without preserving reads.

This work should follow the implementation of #24602 so this fallback can avoid materializing discarded keys when possible.

Built-in implementations can likely do better most of the time. Vector and bitmap-backed accumulators can move or copy only retained entries. Hash tables and collision lists can remap IDs in place or rebuild at retained size. Variable-width and collection accumulators can drop unreferenced buffers and batches. Optimized implementations should update memory accounting and release excess capacity, although exact capacity is not part of the logical API contract.

### Describe alternatives you've considered

- Extending `EmitTo` with arbitrary indices would make its removal and output semantics harder to reason about. Retention does not need to produce aggregate results.
- A free list or tombstone design would avoid immediate compaction, but it would give up dense IDs and require changes throughout vector-backed accumulators. Block-based state may make some removals cheaper in the future, but it does not define how arbitrary survivors are remapped.

### Additional context

- #24602

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.