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..57b46c05859 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,27 @@ 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) + } + } + } + } }