viirya commented on code in PR #26782:
URL: https://github.com/apache/datafusion/pull/26782#discussion_r4418891463
##########
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:
Good catch, confirmed. With `APPROX_DISTINCT(v) GROUP BY k, k2` over input
sorted on `k` only, a 64 KiB limit failed after the emit with "Failed to
allocate additional 64.0 KB ... with 600.0 B already allocated": the table was
empty, but `hashes` was still counted.
The non-clustered `AggregateHashTable::take_state_batch` had the same gap.
It is used by the partial hash stream's early emit and by the single/final hash
streams before spilling. A partial `APPROX_DISTINCT(v) GROUP BY k` over
unsorted input failed the same way under 64 KiB.
In 32b641fb7, both `take_state_batch` now call `shrink_to_fit` on every
accumulator. New tests, each of which fails without its fix:
-
`clustered_partial_stream::tests::emitted_states_release_accumulator_scratch`
(partial group clustering)
- `hash_stream::tests::emitted_states_release_accumulator_scratch` (no
clustering)
##########
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:
Done, thanks. 64 KiB works for both `APPROX_DISTINCT` and `ANY_VALUE`.
Without `GroupsAccumulator::shrink_to_fit`, they fail asking for 58.9 KB and
283.0 KB. The test is now `completed_groups_release_accumulator_scratch`. I
dropped `run_aggregates` and the single-mode test.
##########
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:
Done. `free_scratch_if_empty` and `shrink_to_fit` both call a new
`free_scratch()`.
--
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]
