comphead commented on code in PR #26726:
URL: https://github.com/apache/datafusion/pull/26726#discussion_r4379730364
##########
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:
Design question. The retained capacity sits in the shared
`ClusteredAggregateTable` buffer: `group_indices`, the `GroupValuesPrimitive`
hash table, and the clustered `new_groups` scratch.
`ClusteredFinalAggregateStream` and `ClusteredSingleAggregateStream` emit
through the same `EmitTo::First` path and cannot emit the active group. So I
expect them to fail the same way under a limit below one batch of capacity,
even after #26637. #26637 already hit this retention under partial clustering.
Releasing that capacity on the refusal path after
`take_completed_state_batch` would cover all three streams with one change. For
example, call `group_indices.shrink_to_fit()` and add a shrink hook on the
group values. It would also keep one state row per group and match the legacy
`GroupedHashAggregateStream`, which never emits the active group under full
clustering. Was that considered? If the reset is preferred for `Partial`, it
would help to say in the PR that Final and Single need a different fix.
##########
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:
This shows the partial stream no longer has an error path for a non-empty
batch under full clustering. It finishes with one group under a 6,000-byte
limit. The legacy stream returns the refusal here, because it only emits early
when `group_values.len() > 1`. If that is intended, the name
`full_clustering_refusal_is_an_error` no longer fits the partial case.
Splitting the partial case into its own test, or renaming it, would keep the
name accurate.
##########
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:
This duplicates `memory_pressure_releases_retained_capacity`. Both use 8192
distinct keys per batch, one and two `Int64` keys, the 100 KiB limit, and
compare the merged results. The only extra check is that the planner picks
`mode=Partial` with `group_clustering_mode=Full`, and `.slt` plans already show
that (`group_by.slt`, `agg_func_substitute.slt`). Dropping this file also
removes `PartialStreamLimitPool` and the second batch builder.
If it stays, `run` plans the query twice: once in `create_physical_plan` and
again in `df.collect()`. Collecting the plan that was already built avoids the
second planning pass.
##########
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:
This test and `completed_groups_under_memory_pressure` now go through the
same branch. If the ranges here overlap by one key (for example `0..8193,
8192..16385, 16384..24576`), this test also covers the split state rows, under
a limit that fails on `main`. You could then assert the extra rows, drop the
older test, and drop the `held` sampling from `run`.
--
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]
