viirya commented on code in PR #26726:
URL: https://github.com/apache/datafusion/pull/26726#discussion_r4383721472
##########
datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs:
##########
@@ -319,13 +321,10 @@ impl Aggregating {
let output = match
reservation.try_resize(self.table.memory_size()) {
Ok(()) => self.table.take_completed_state_batch()?,
Err(oom @ DataFusionError::ResourcesExhausted(_)) => {
- // Partial group clustering may have an unbounded active
key range.
- // The final stage can merge incomplete states emitted
here.
- if matches!(self.table.group_clustering(),
GroupClustering::Full(_)) {
-
Some(self.table.take_completed_state_batch()?.ok_or(oom)?)
- } else {
- Some(self.table.take_state_batch()?.ok_or(oom)?)
- }
+ // Emit every state, including the active groups, and
reset the
+ // table to release the capacity it grew for this batch.
The
+ // final stage merges the states of a group split across
outputs.
+ Some(self.table.take_state_batch()?.ok_or(oom)?)
Review Comment:
Yes, switched to this in e7403778e. On a refusal under full clustering, the
stream emits the completed groups and then calls a new
`ClusteredAggregateTable::shrink_to_fit`. It clears and shrinks `group_indices`
and calls a new `GroupValues::shrink_to_fit`, which does nothing by default.
The implementations:
- `GroupValuesPrimitive`: the hash table and `values`.
- `GroupValuesClustered`: `new_groups`.
- `GroupValuesColumn`: the hash table and `hashes_buffer`. In the
non-streaming case, also the emit and vectorized scratch.
- `GroupValuesRows`: the hash table, `hashes_buffer`, and `rows_buffer`.
`GroupValuesBytes` and `GroupValuesBytesView` rebuild their map on
`EmitTo::First`, so they need nothing.
The stream emits one state row per group again, and it returns the refusal
when no group is complete, as the legacy stream does. Final and Single can call
the same method after their completed-groups emit.
##########
datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs:
##########
@@ -822,8 +824,10 @@ async fn full_clustering_refusal_is_an_error() ->
Result<()> {
ClusteredPartialAggregateStream[0] try_grow [bytes]
ClusteredPartialAggregateStream[0] try_grow [bytes] -> refused
ClusteredPartialAggregateStream[0] shrink [bytes]
+ ClusteredPartialAggregateStream[0] shrink [bytes]
+ ClusteredPartialAggregateStream[0] shrink [bytes]
ClusteredPartialAggregateStream[0] unregister
- => Err(Resources exhausted: Failed to allocate additional [bytes] for
ClusteredPartialAggregateStream[0] with [bytes] already allocated for this
reservation - [bytes] remain available for the total memory pool: greedy(used:
[bytes], pool_size: [bytes]))
+ => Ok(1 rows in 1 batches)
Review Comment:
Since the stream again emits only the completed groups, the partial case
returns the refusal here again. I restored the original snapshot, so the name
fits again.
##########
datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs:
##########
@@ -487,13 +512,63 @@ mod tests {
Ok((output, held))
}
+ /// Merges the partial `SUM` states by key, as the final stage does.
+ fn merge(output: &[RecordBatch]) -> BTreeMap<Vec<i64>, i64> {
+ let mut merged = BTreeMap::new();
+ for batch in output {
+ let columns: Vec<_> = batch
+ .columns()
+ .iter()
+ .map(|column| column.as_primitive::<Int64Type>())
+ .collect();
+ let (sum, keys) = columns.split_last().unwrap();
+ for row in 0..batch.num_rows() {
+ let key = keys.iter().map(|k| k.value(row)).collect();
+ *merged.entry(key).or_default() += sum.value(row);
+ }
+ }
+ merged
+ }
+
+ fn num_rows(output: &[RecordBatch]) -> usize {
+ output.iter().map(|batch| batch.num_rows()).sum()
+ }
+
/// The limit fits the table with only the active group, but not the table
- /// before the completed groups are emitted. The result must not change.
+ /// before the completed groups are emitted. Key 256 spans both batches, so
+ /// emitting its state early splits it into two state rows, which the final
+ /// stage merges.
#[tokio::test]
async fn completed_groups_under_memory_pressure() -> Result<()> {
- let (expected, held) =
run(Arc::new(UnboundedMemoryPool::default())).await?;
- let (actual, _) = run(Arc::new(GreedyMemoryPool::new(held))).await?;
- assert_eq!(batches_to_string(&actual), batches_to_string(&expected));
+ let input = || batches(1, &[0..257, 256..513]);
+ let (expected, held) =
+ run(1, input()?, Arc::new(UnboundedMemoryPool::default())).await?;
+ let (actual, _) = run(1, input()?,
Arc::new(GreedyMemoryPool::new(held))).await?;
+ assert_eq!(num_rows(&actual), num_rows(&expected) + 1);
+ assert_eq!(merge(&actual), merge(&expected));
+ Ok(())
+ }
+
+ /// Every row is a distinct group, so a batch grows the table for 8192
+ /// groups: the single-key hash table (~272 KB), the clustered `new_groups`
+ /// scratch and `group_indices` (64 KB each). Emitting only the completed
+ /// groups would keep that capacity reserved and fail this limit, although
+ /// the active group needs a few KB.
+ #[tokio::test]
+ async fn memory_pressure_releases_retained_capacity() -> Result<()> {
Review Comment:
Done. `completed_groups_under_memory_pressure` now uses the overlapping
ranges, so the active group spans each batch boundary. It runs under the 100
KiB limit, which fails on `main`. It covers the four group-values
implementations that implement the new hook: `Int64`, `Int64, Int64`, `Float64,
Float64`, and `Int64, List`. Without the hook, the `Int64`, `Float64, Float64`,
and `Int64, List` cases still fail when only `group_indices` is released.
Because the output no longer changes, the test compares it exactly. The older
test and the `held` sampling are removed.
##########
datafusion/core/tests/memory_limit/clustered_partial_oom.rs:
##########
@@ -0,0 +1,171 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A partial aggregate over input sorted on every grouping key must finish
+//! when its table needs only a few KB between batches.
+//!
+//! On a refused reservation the clustered partial stream emits its states and
+//! resets the table. Emitting only the completed groups would leave one active
+//! group but keep the capacity the table grew for the whole batch, so
+//! re-reserving it would fail although the active group alone fits easily.
+//!
+//! The downstream merge and final aggregate cannot spill and need at least
+//! a batch of memory themselves, so a shared limit cannot single out the
+//! partial stream. The pool here limits only the partial streams.
+
+use std::fmt::{self, Display};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use arrow::array::{ArrayRef, Int64Array, RecordBatch};
+use arrow::datatypes::{DataType, Field, Schema};
+use datafusion::datasource::MemTable;
+use datafusion::physical_plan::displayable;
+use datafusion::prelude::*;
+use datafusion_common::test_util::batches_to_sort_string;
+use datafusion_common::{Result, resources_datafusion_err};
+use datafusion_execution::memory_pool::{
+ MemoryPool, MemoryReservation, UnboundedMemoryPool,
+};
+use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+
+const ROWS: usize = 8192;
+const PARTITIONS: usize = 4;
+const PARTIAL_STREAM: &str = "ClusteredPartialAggregateStream";
+
+/// Unbounded pool except that each partial aggregate stream may reserve at
+/// most `limit` bytes.
+#[derive(Debug)]
+struct PartialStreamLimitPool {
+ limit: usize,
+ used: AtomicUsize,
+}
+
+impl Display for PartialStreamLimitPool {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ write!(f, "{}(limit: {})", self.name(), self.limit)
+ }
+}
+
+impl MemoryPool for PartialStreamLimitPool {
+ fn name(&self) -> &str {
+ "partial_stream_limit"
+ }
+
+ fn grow(&self, _reservation: &MemoryReservation, additional: usize) {
+ self.used.fetch_add(additional, Ordering::Relaxed);
+ }
+
+ fn shrink(&self, _reservation: &MemoryReservation, shrink: usize) {
+ self.used.fetch_sub(shrink, Ordering::Relaxed);
+ }
+
+ fn try_grow(&self, reservation: &MemoryReservation, additional: usize) ->
Result<()> {
+ if reservation.consumer().name().starts_with(PARTIAL_STREAM)
+ && reservation.size() + additional > self.limit
+ {
+ return Err(resources_datafusion_err!(
+ "Failed to allocate additional {additional} bytes for {} with
{} bytes \
+ already allocated, limit {}",
+ reservation.consumer().name(),
+ reservation.size(),
+ self.limit
+ ));
+ }
+ self.grow(reservation, additional);
+ Ok(())
+ }
+
+ fn reserved(&self) -> usize {
+ self.used.load(Ordering::Relaxed)
+ }
+}
+
+/// Each partition holds three batches sorted on every key, and every row is
+/// a distinct group.
+fn table(num_keys: usize) -> Result<MemTable> {
+ let mut fields: Vec<_> = (0..num_keys)
+ .map(|i| Field::new(format!("k{i}"), DataType::Int64, false))
+ .collect();
+ fields.push(Field::new("v", DataType::Int64, false));
+ let schema = Arc::new(Schema::new(fields));
+ let partitions = (0..PARTITIONS)
+ .map(|p| {
+ (0..3)
+ .map(|b| {
+ let start = ((p * 3 + b) * ROWS) as i64;
+ let ids = start..start + ROWS as i64;
+ let columns = (0..=num_keys)
+ .map(|_| {
+ Arc::new(Int64Array::from_iter_values(ids.clone()))
+ as ArrayRef
+ })
+ .collect();
+ Ok(RecordBatch::try_new(Arc::clone(&schema), columns)?)
+ })
+ .collect::<Result<Vec<_>>>()
+ })
+ .collect::<Result<Vec<_>>>()?;
+ let order = (0..num_keys)
+ .map(|i| col(format!("k{i}")).sort(true, false))
+ .collect();
+ Ok(MemTable::try_new(schema, partitions)?.with_sort_order(vec![order]))
+}
+
+async fn run(num_keys: usize, pool: Arc<dyn MemoryPool>) ->
Result<Vec<RecordBatch>> {
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_pool(pool)
+ .build_arc()?;
+ let config = SessionConfig::new().with_target_partitions(PARTITIONS);
+ let ctx = SessionContext::new_with_config_rt(config, runtime);
+ ctx.register_table("t", Arc::new(table(num_keys)?))?;
+ let keys = (0..num_keys)
+ .map(|i| format!("k{i}"))
+ .collect::<Vec<_>>()
+ .join(", ");
+ let df = ctx
+ .sql(&format!("select {keys}, sum(v) from t group by {keys}"))
+ .await?;
+ let plan = df.clone().create_physical_plan().await?;
+ let plan = displayable(plan.as_ref()).indent(true).to_string();
+ assert!(
+ plan.contains("mode=Partial") &&
plan.contains("group_clustering_mode=Full"),
+ "{plan}"
+ );
+ df.collect().await
+}
+
+/// With 8192 distinct keys per batch, the capacity a table grows for a batch
+/// is about 344 KB for a single key and 131 KB for two keys, above the limit.
+#[tokio::test]
+async fn sorted_partial_aggregate_completes_under_memory_pressure() ->
Result<()> {
Review Comment:
Agreed, removed.
##########
datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs:
##########
@@ -319,13 +321,10 @@ impl Aggregating {
let output = match
reservation.try_resize(self.table.memory_size()) {
Ok(()) => self.table.take_completed_state_batch()?,
Err(oom @ DataFusionError::ResourcesExhausted(_)) => {
- // Partial group clustering may have an unbounded active
key range.
- // The final stage can merge incomplete states emitted
here.
- if matches!(self.table.group_clustering(),
GroupClustering::Full(_)) {
-
Some(self.table.take_completed_state_batch()?.ok_or(oom)?)
- } else {
- Some(self.table.take_state_batch()?.ok_or(oom)?)
- }
+ // Emit every state, including the active groups, and
reset the
+ // table to release the capacity it grew for this batch.
The
+ // final stage merges the states of a group split across
outputs.
+ Some(self.table.take_state_batch()?.ok_or(oom)?)
Review Comment:
Thanks, confirmed. With your inputs and the 13,000-byte limit, a
`SortExec(k, sum)` above the partial passes `(256, 9), (256, 1)` through
unsorted. The stream no longer resets the table under full clustering. It emits
the completed groups and releases the retained capacity instead, so the active
group stays in one row, and `Unique(k)` holds within a partition again.
`completed_groups_under_memory_pressure` compares the output exactly under the
limit, so it fails if a group is split again.
##########
datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs:
##########
@@ -319,13 +321,10 @@ impl Aggregating {
let output = match
reservation.try_resize(self.table.memory_size()) {
Ok(()) => self.table.take_completed_state_batch()?,
Err(oom @ DataFusionError::ResourcesExhausted(_)) => {
- // Partial group clustering may have an unbounded active
key range.
- // The final stage can merge incomplete states emitted
here.
- if matches!(self.table.group_clustering(),
GroupClustering::Full(_)) {
-
Some(self.table.take_completed_state_batch()?.ok_or(oom)?)
- } else {
- Some(self.table.take_state_batch()?.ok_or(oom)?)
- }
+ // Emit every state, including the active groups, and
reset the
+ // table to release the capacity it grew for this batch.
The
Review Comment:
Confirmed. `HllGroupsAccumulator::size` counts the `hashes` scratch, and
nothing releases it after `state`. That memory belongs to the accumulator
rather than the table buffer. Releasing it needs a shrink hook on
`GroupsAccumulator`, which reaches every accumulator that keeps per-batch
scratch. I'd like to keep this PR to the group values, and I opened #26781 for
the accumulators.
--
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]
