From ee5d84128c24e99df37546d348a28926fc027910 Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 12 Oct 2026 13:23:25 +0100 Subject: [PATCH 1/2] perf: drain a large ordered aggregate range once instead of in batch-size steps An aggregate over input sorted on a prefix of its grouping keys runs in DataFusion's ordered partial or ordered final stream. Groups become ready only when the sorted prefix changes or the input ends, and the ready range was then emitted with EmitTo::First(batch_size) per output batch. Each such emit walks the whole group hash table and copies every remaining key and state, so a range of N groups took O(N^2 / batch_size). COUNT(DISTINCT) after a window partitioned by the same prefix hits this: on push_ranking_results one window partition held ~30M groups and its Partial and PartialMerge aggregates took ~830 s each where Spark needs 89 s for the whole query. When a ready range is larger than batch_size and the stream outputs states (Partial and PartialReduce), emit the range once into a materialized batch and hand it out in batch_size slices, as the unordered hash streams already do. The undelivered slices count towards the table's memory and emptiness, and are prepended to the state batch taken under memory pressure. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../aggregates/aggregate_hash_table/common.rs | 11 + .../aggregate_hash_table/common_ordered.rs | 58 +++- .../src/aggregates/mod.rs | 2 + .../src/aggregates/ordered_drain_tests.rs | 322 ++++++++++++++++++ .../comet/exec/CometAggregateSuite.scala | 25 ++ 5 files changed, 416 insertions(+), 2 deletions(-) create mode 100644 native/vendor/datafusion-physical-plan/src/aggregates/ordered_drain_tests.rs diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs index 5c6521cc354..85327d7a988 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -512,6 +512,17 @@ impl MaterializedAggregateOutput { self.offset >= self.batch.num_rows() } + pub(super) fn take_rest(&mut self) -> Option { + if self.is_exhausted() { + return None; + } + let rest = self + .batch + .slice(self.offset, self.batch.num_rows() - self.offset); + self.offset = self.batch.num_rows(); + Some(rest) + } + pub(super) fn memory_size(&self) -> usize { self.batch.get_array_memory_size() } diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index 2293e7b1b8e..0e604fee8c1 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -21,6 +21,7 @@ use std::marker::PhantomData; use std::sync::Arc; +use arrow::compute::concat_batches; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; @@ -38,7 +39,9 @@ use crate::aggregates::{ evaluate_group_by, }; -use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; +use super::common::{ + AggregateAccumulator, EvaluatedAggregateBatch, MaterializedAggregateOutput, +}; /// Aggregate table shared by the ordered partial and final paths. /// @@ -94,6 +97,8 @@ pub(in crate::aggregates) struct OrderedAggregateTable { /// Group keys, ordering state, and accumulator states. pub(super) buffer: OrderedAggregateTableBuffer, + pending: Option, + _mode: PhantomData, } @@ -183,6 +188,7 @@ impl OrderedAggregateTable { group_indices: vec![], accumulators, }, + pending: None, _mode: PhantomData, }) } @@ -236,6 +242,10 @@ impl OrderedAggregateTable { /// Check if there is zero groups accumulated so far. pub(in crate::aggregates) fn is_empty(&self) -> bool { self.num_groups() == 0 + && self + .pending + .as_ref() + .is_none_or(MaterializedAggregateOutput::is_exhausted) } /// All internal buffer's memory size. @@ -248,6 +258,10 @@ impl OrderedAggregateTable { + self.buffer.group_values.size() + self.buffer.group_ordering.size() + self.buffer.group_indices.allocated_size() + + self + .pending + .as_ref() + .map_or(0, MaterializedAggregateOutput::memory_size) } pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics { @@ -264,8 +278,15 @@ impl OrderedAggregateTable { pub(in crate::aggregates) fn take_state_batch( &mut self, ) -> Result> { + let rest = match self.pending.take().and_then(|mut p| p.take_rest()) { + Some(rest) => Some(RecordBatch::try_new( + Arc::clone(&self.state_schema), + rest.columns().to_vec(), + )?), + None => None, + }; if self.buffer.group_values.is_empty() { - return Ok(None); + return Ok(rest); } let mut output = self.buffer.group_values.emit(EmitTo::All)?; @@ -275,6 +296,10 @@ impl OrderedAggregateTable { let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?; debug_assert!(batch.num_rows() > 0); + let batch = match rest { + Some(rest) => concat_batches(&self.state_schema, &[rest, batch])?, + None => batch, + }; // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the // key/index buffers too so the memory reservation can be released @@ -371,6 +396,13 @@ impl OrderedAggregateTable { &mut self, is_final: bool, ) -> Result> { + if let Some(pending) = &mut self.pending { + if let Some(batch) = pending.next_batch(self.batch_size) { + return Ok(Some(batch)); + } + self.pending = None; + } + if self.buffer.group_values.is_empty() { return Ok(None); } @@ -378,6 +410,28 @@ impl OrderedAggregateTable { let Some(emit_to) = self.buffer.group_ordering.emit_to() else { return Ok(None); }; + + let ready = match emit_to { + EmitTo::First(n) => n, + EmitTo::All => self.buffer.group_values.len(), + }; + if !is_final && ready > self.batch_size { + let timer = self.group_by_metrics.emitting_time.timer(); + let mut output = self.buffer.group_values.emit(emit_to)?; + if let EmitTo::First(n) = emit_to { + self.buffer.group_ordering.remove_groups(n); + } + for acc in &mut self.buffer.accumulators { + output.extend(acc.state(emit_to)?); + } + drop(timer); + let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + let mut pending = MaterializedAggregateOutput::new(batch); + let first = pending.next_batch(self.batch_size); + self.pending = Some(pending); + return Ok(first); + } + let (emit_to, should_remove_groups) = self.clamp_emit_to(self.buffer.group_values.len(), emit_to); diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs index e69d1bad404..321b895f6d0 100644 --- a/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs +++ b/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs @@ -218,6 +218,8 @@ mod single_stream; mod skip_partial; // COMET PATCH: tests for an aggregate whose first reservation fails and spills. #[cfg(test)] +mod ordered_drain_tests; +#[cfg(test)] mod starved_spill_tests; mod topk; diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/ordered_drain_tests.rs b/native/vendor/datafusion-physical-plan/src/aggregates/ordered_drain_tests.rs new file mode 100644 index 00000000000..b5d029997ef --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/ordered_drain_tests.rs @@ -0,0 +1,322 @@ +// 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. + +use std::collections::HashSet; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use arrow::array::{Array, Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion_common::Result; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_functions_aggregate::count::count_udaf; +use datafusion_physical_expr::aggregate::AggregateExprBuilder; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; + +use super::{AggregateExec, AggregateMode, PhysicalGroupBy, StreamType}; +use crate::common::collect; +use crate::execution_plan::ExecutionPlan; +use crate::test::TestMemoryExec; + +const BATCH_SIZE: usize = 256; + +fn raw_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("p", DataType::Utf8, false), + Field::new("s", DataType::Utf8, false), + Field::new("v", DataType::Int64, false), + ])) +} + +fn raw_batches( + schema: &SchemaRef, + prefixes: &[&str], + rows_per_prefix: usize, +) -> Result> { + let mut batches = vec![]; + for (i, prefix) in prefixes.iter().enumerate() { + for start in (0..rows_per_prefix).step_by(4096) { + let end = (start + 4096).min(rows_per_prefix); + let rows = start..end; + batches.push(RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(StringArray::from_iter_values( + rows.clone().map(|_| *prefix), + )), + Arc::new(StringArray::from_iter_values(rows.clone().map(|r| { + format!("ranking-{i}-{:012}", (r * 7919) % rows_per_prefix) + }))), + Arc::new(Int64Array::from_iter_values(rows.map(|r| r as i64))), + ], + )?); + } + } + Ok(batches) +} + +fn aggregate( + mode: AggregateMode, + input_schema: &SchemaRef, + raw_schema: &SchemaRef, + input: Arc, +) -> Result> { + let group_by = PhysicalGroupBy::new_single(vec![ + (col("p", input_schema)?, "p".to_string()), + (col("s", input_schema)?, "s".to_string()), + ]); + let count = Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("v", raw_schema)?]) + .schema(Arc::clone(raw_schema)) + .alias("count_v") + .build()?, + ); + Ok(Arc::new(AggregateExec::try_new( + mode, + group_by, + vec![count], + vec![None], + input, + Arc::clone(raw_schema), + )?)) +} + +fn input( + batches: &[RecordBatch], + schema: &SchemaRef, + sorted: bool, +) -> Result> { + let ordering = if sorted { + vec![ + LexOrdering::new(vec![PhysicalSortExpr::new_default(col("p", schema)?)]) + .unwrap(), + ] + } else { + vec![] + }; + let exec = TestMemoryExec::try_new(&[batches.to_vec()], Arc::clone(schema), None)? + .try_with_sort_information(ordering)?; + Ok(Arc::new(TestMemoryExec::update_cache(&Arc::new(exec)))) +} + +fn task_ctx(pool: Option>) -> Result> { + let mut runtime = RuntimeEnvBuilder::new(); + if let Some(pool) = pool { + runtime = runtime.with_memory_pool(pool); + } + Ok(Arc::new( + TaskContext::default() + .with_session_config( + SessionConfig::new() + .with_batch_size(BATCH_SIZE) + .set_bool("datafusion.execution.enable_migration_aggregate", true) + .set_str( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + "1.1", + ), + ) + .with_runtime(runtime.build_arc()?), + )) +} + +struct Run { + elapsed: Duration, + groups: HashSet<(String, String)>, + total_count: i64, + rows: usize, +} + +async fn run( + aggregate: Arc, + pool: Option>, + expected_stream: fn(&StreamType) -> bool, + bounded_batches: bool, +) -> Result { + let ctx = task_ctx(pool)?; + assert!(expected_stream(&aggregate.execute_typed(0, &ctx)?)); + let start = Instant::now(); + let output = collect(aggregate.execute(0, ctx)?).await?; + let elapsed = start.elapsed(); + let mut groups = HashSet::new(); + let mut total_count = 0; + let mut rows = 0; + for batch in &output { + assert!( + !bounded_batches || batch.num_rows() <= BATCH_SIZE, + "batch of {} rows", + batch.num_rows() + ); + rows += batch.num_rows(); + let p = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let s = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let c = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + groups.insert((p.value(i).to_string(), s.value(i).to_string())); + total_count += c.value(i); + } + } + Ok(Run { + elapsed, + groups, + total_count, + rows, + }) +} + +fn is_ordered_partial(stream: &StreamType) -> bool { + matches!(stream, StreamType::OrderedPartialAggregate(_)) +} + +fn is_partial_hash(stream: &StreamType) -> bool { + matches!(stream, StreamType::PartialHash(_)) +} + +fn is_ordered_final(stream: &StreamType) -> bool { + matches!(stream, StreamType::OrderedFinalAggregate(_)) +} + +fn is_any(_: &StreamType) -> bool { + true +} + +#[tokio::test] +async fn ordered_partial_aggregate_drains_a_huge_ordered_range_in_linear_time() +-> Result<()> { + let schema = raw_schema(); + let rows = 300_000; + let batches = raw_batches(&schema, &["x"], rows)?; + + let sorted = run( + aggregate( + AggregateMode::Partial, + &schema, + &schema, + input(&batches, &schema, true)?, + )?, + None, + is_ordered_partial, + true, + ) + .await?; + let unsorted = run( + aggregate( + AggregateMode::Partial, + &schema, + &schema, + input(&batches, &schema, false)?, + )?, + None, + is_partial_hash, + true, + ) + .await?; + + assert_eq!(sorted.rows, rows); + assert_eq!(sorted.groups.len(), rows); + assert_eq!(sorted.total_count, rows as i64); + assert_eq!(sorted.groups, unsorted.groups); + assert!( + sorted.elapsed < unsorted.elapsed * 5 + Duration::from_millis(500), + "ordered {:?}, unordered {:?}", + sorted.elapsed, + unsorted.elapsed + ); + Ok(()) +} + +#[tokio::test] +async fn ordered_partial_reduce_drains_a_huge_ordered_range_in_linear_time() -> Result<()> +{ + let schema = raw_schema(); + let rows = 300_000; + let batches = raw_batches(&schema, &["x"], rows)?; + let partial = aggregate( + AggregateMode::Partial, + &schema, + &schema, + input(&batches, &schema, false)?, + )?; + let states = collect(partial.execute(0, task_ctx(None)?)?).await?; + let state_schema = partial.schema(); + + let reduce = |sorted: bool| -> Result> { + aggregate( + AggregateMode::PartialReduce, + &state_schema, + &schema, + input(&states, &state_schema, sorted)?, + ) + }; + let pool = || -> Option> { + Some(Arc::new(GreedyMemoryPool::new(1 << 30))) + }; + let sorted = run(reduce(true)?, pool(), is_ordered_final, true).await?; + let unsorted = run(reduce(false)?, pool(), is_any, true).await?; + + assert_eq!(sorted.rows, rows); + assert_eq!(sorted.total_count, rows as i64); + assert_eq!(sorted.groups, unsorted.groups); + assert!( + sorted.elapsed < unsorted.elapsed * 5 + Duration::from_millis(500), + "ordered {:?}, unordered {:?}", + sorted.elapsed, + unsorted.elapsed + ); + Ok(()) +} + +#[tokio::test] +async fn ordered_partial_aggregate_keeps_drained_groups_when_memory_runs_out() +-> Result<()> { + let schema = raw_schema(); + let rows = 60_000; + let batches = raw_batches(&schema, &["a", "b", "c", "d"], rows)?; + let pool: Arc = Arc::new(GreedyMemoryPool::new(6 * 1024 * 1024)); + let constrained = run( + aggregate( + AggregateMode::Partial, + &schema, + &schema, + input(&batches, &schema, true)?, + )?, + Some(Arc::clone(&pool)), + is_ordered_partial, + false, + ) + .await?; + + assert_eq!(constrained.total_count, 4 * rows as i64); + assert_eq!(constrained.groups.len(), 4 * rows); + assert_eq!(pool.reserved(), 0); + Ok(()) +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index eb2819dcd99..288439b13ec 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -3728,4 +3728,29 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { assert(sortAggregates(cometPlan).nonEmpty, s"plan:\n$cometPlan") } } + + test("count distinct over a window-sorted input with one huge ordered key range") { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "2", + CometConf.COMET_BATCH_SIZE.key -> "64") { + withTempPath { path => + spark + .range(0, 60000, 1, 4) + .selectExpr( + "CAST(DATE'2026-10-10' + CAST(id % 2 AS INT) AS DATE) AS d", + "IF(id % 50 = 0, 'b', 'a') AS t", + "concat('r', CAST((id * 7919) % 30000 AS STRING)) AS r", + "CAST(id % 13 AS DOUBLE) AS w") + .write + .parquet(path.getAbsolutePath) + withParquetTable(path.getAbsolutePath, "ranked") { + val df = sql(""" + |SELECT d, t, ROUND(MAX(mw), 3) AS mw, COUNT(DISTINCT r) AS rankings + |FROM (SELECT d, t, r, MAX(w) OVER (PARTITION BY d, t) AS mw FROM ranked) + |GROUP BY d, t""".stripMargin) + checkSparkAnswerAndOperator(df) + } + } + } + } } From 3ef370e2555e9ef4117b44bc1aebeb9b084a47ea Mon Sep 17 00:00:00 2001 From: msaf Date: Mon, 12 Oct 2026 13:36:19 +0100 Subject: [PATCH 2/2] style: format the new aggregate test for spotless Co-Authored-By: Claude Opus 5.5 (1M context) --- .../scala/org/apache/comet/exec/CometAggregateSuite.scala | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 288439b13ec..57b46c05859 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -3730,9 +3730,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } test("count distinct over a window-sorted input with one huge ordered key range") { - withSQLConf( - SQLConf.SHUFFLE_PARTITIONS.key -> "2", - CometConf.COMET_BATCH_SIZE.key -> "64") { + withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2", CometConf.COMET_BATCH_SIZE.key -> "64") { withTempPath { path => spark .range(0, 60000, 1, 4)