Repository navigation
Description
Describe the bug
#26726 (for #26725) releases the capacity that a clustered aggregate table keeps after it emits its completed groups under memory pressure, but only for the group values and group_indices. Accumulators can also keep per-batch scratch, and GroupsAccumulator has no way to release it.
HllGroupsAccumulator (APPROX_DISTINCT) keeps a hashes: Vec<u64> sized to the largest input batch. size() counts it, and state does not release it:
datafusion/datafusion/functions-aggregate/src/approx_distinct.rs
Lines 593 to 605 in c18e5d9
| fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { | |
| let mut groups = emit_to.take_needed(&mut self.groups); | |
| let mut builder = BinaryBuilder::new(); | |
| let mut scratch: Vec<u8> = Vec::new(); | |
| let mut freed = 0; | |
| for g in groups.iter_mut() { | |
| freed += g.heap_bytes(); | |
| g.serialize(&mut scratch); | |
| builder.append_value(&scratch); | |
| } | |
| // The emitted groups have been removed; reclaim their tracked bytes. | |
| self.allocated_bytes = self.allocated_bytes.saturating_sub(freed); | |
| Ok(vec![Arc::new(builder.finish())]) |
datafusion/datafusion/functions-aggregate/src/approx_distinct.rs
Lines 640 to 644 in c18e5d9
| fn size(&self) -> usize { | |
| self.groups.capacity() * size_of::<GroupHll>() | |
| + self.allocated_bytes | |
| + self.hashes.capacity() * size_of::<u64>() | |
| } |
With 8192-row batches, each APPROX_DISTINCT keeps 64 KiB after the emit. As reported by @sunchao in #26726, a fully clustered partial aggregate with APPROX_DISTINCT(k) and APPROX_DISTINCT(v) over one sorted batch of 8192 distinct Int64 keys still fails a 100 KiB limit, although the table holds only the active group.
To Reproduce
A partial APPROX_DISTINCT(k), APPROX_DISTINCT(v) GROUP BY k over input sorted on k, with 8192 distinct keys per batch and a 100 KiB GreedyMemoryPool.
Expected behavior
After the memory-pressure emit, the reservation should cover only the active group's state, so the stream continues.
Additional context
One option: add a GroupsAccumulator::shrink_to_fit that does nothing by default, call it from ClusteredAggregateTable::shrink_to_fit, and implement it for accumulators that keep per-batch scratch, starting with HllGroupsAccumulator.
Willingness to contribute
I plan to contribute a fix for this bug independently
Activity
fix: release clustered aggregate table capacity after emitting comple…
fix: release per-batch accumulator scratch after a clustered partial …
fix: release per-batch accumulator scratch after a clustered aggregat…