comphead commented on code in PR #26782:
URL: https://github.com/apache/datafusion/pull/26782#discussion_r4417887134
##########
datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs:
##########
@@ -653,6 +653,14 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
self.allocation_bytes
}
+ fn shrink_to_fit(&mut self) {
+ let vec_size_pre = self.states.allocated_size();
+ self.states.shrink_to_fit();
+ self.adjust_allocation(vec_size_pre, self.states.allocated_size());
+ let scratch = std::mem::take(&mut self.scratch);
+ self.free_allocation(scratch.allocated_size());
+ }
Review Comment:
`free_scratch_if_empty` already has the take-and-free logic. Could both call
a small `free_scratch()` helper, so the scratch accounting lives in one place?
##########
datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs:
##########
@@ -559,4 +574,66 @@ mod tests {
assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full);
assert_completed_groups_emitted_under_limit(&plan).await
}
+
+ /// Like [`completed_groups_under_memory_pressure`], but with accumulators
+ /// that keep per-batch scratch. For example, each `APPROX_DISTINCT` keeps
+ /// the hashes of the last batch (64 KB for 8192 rows), and the
+ /// `GroupsAccumulatorAdapter` of `ANY_VALUE` keeps its per-batch indices
+ /// (~160 KB). After the completed groups are emitted, the accumulators
must
+ /// release that scratch too, or reserving the table fails the limit
+ /// although the active group needs a few KB. The output must not change.
+ #[tokio::test]
+ async fn accumulator_scratch_under_memory_pressure() -> Result<()> {
Review Comment:
This could reuse the shared helpers that
`completed_groups_release_accumulator_scratch` already uses. That drops the
`run_aggregates` split, the local `build`, the `Aggregates` alias and three
imports:
```rust
for udaf in [approx_distinct_udaf(), any_value_udaf()] {
let plan = aggregate_over_sorted_input(
AggregateMode::Partial,
&["k"],
distinct_key_batches(),
udaf,
)?;
assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full);
assert_output_unchanged_under_limit(&plan, 64 * 1024).await?;
}
```
It builds one `APPROX_DISTINCT` instead of two, so the limit has to sit
below a single `hashes` buffer, as in the single-mode test. I expect 64 KiB to
work for both cases but haven't run it.
All three streams go through the same
`ClusteredAggregateTable::shrink_to_fit`, and the single stream's call site is
already covered by `completed_groups_release_batch_capacity`. So with the loop
above, `completed_groups_release_accumulator_scratch` repeats the same path and
one of the two tests could go.
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_clustered.rs:
##########
@@ -287,14 +287,16 @@ impl<AggrMode> ClusteredAggregateTable<AggrMode> {
}
/// Releases the capacity that the buffer holds beyond its stored groups,
- /// such as hash table slots and the per-batch `group_indices` scratch, so
- /// that [`Self::memory_size`] reflects the remaining groups.
- ///
- /// Accumulator capacity is not released.
+ /// such as hash table slots, the per-batch `group_indices` scratch and
+ /// per-batch accumulator scratch, so that [`Self::memory_size`] reflects
+ /// the remaining groups.
pub(in crate::aggregates) fn shrink_to_fit(&mut self) {
self.buffer.group_values.shrink_to_fit();
self.buffer.group_indices.clear();
self.buffer.group_indices.shrink_to_fit();
+ for acc in &mut self.buffer.accumulators {
+ acc.shrink_to_fit();
Review Comment:
`take_state_batch` below has the same gap. The partial stream calls it in
the OOM branch for partial group clustering. It clears the key and index
buffers but leaves `HllGroupsAccumulator::hashes` counted, and `Outputting`
then reserves `memory_size()` again. From reading the code (I haven't run it),
`APPROX_DISTINCT(k0), APPROX_DISTINCT(v) GROUP BY k0, k1` with the input sorted
on `k0` only should still fail under the 100 KiB pool, because the two `hashes`
buffers alone are about 128 KB. Could `take_state_batch` run the same
accumulator loop?
--
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]
