From 296b4c497b41a7c6943a9200e62163b0bca545dc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 28 Oct 2026 02:55:33 +0000 Subject: [PATCH] fix: [branch-1.2] spill a native PartialMerge whose input is sorted on some of its keys (#7222) * fix: spill a native PartialMerge whose input is sorted on some of its keys Comet runs a Spark PartialMerge as DataFusion Single so that it spills rather than emit a key twice (#6857). When the input is sorted on a prefix of the grouping keys, DataFusion 55.2 runs Single with the legacy GroupedHashAggregateStream, whose spill asks for more than its share of the pool to sort the batch it writes, and the task fails. Plan that PartialMerge as DataFusion Final instead. Final runs an OrderedFinalAggregateStream there, which spills and reads its spill files back, and MergeAsPartial evaluates to its state under Final as it does under Single. Final reads each state from the columns after the keys and spills its groups in the input's types, so the planner switches only when the input is the grouping keys followed by native-typed states. A Spark partial's binary collect buffers keep the Single plan. Closes #6880. * refactor: decide the PartialMerge aggregate mode before building it Read the input order from the child's equivalence properties, as AggregateExec does, instead of building the aggregate as Single and again as Final. Simplify the state type guard, note when the switch can be rechecked, and cover input sorted on every key in the planner test. * address review: lazy sort check, document fair-limit exemption, cover changing g and a sort-merge join * address review: walk state types once, share planner test helpers * test: cover the PartialMerge inputs Final rejects and run the join case under AQE Add planner cases for swapped keys, states after one key and a column after the states, which each keep the PartialMerge as Single. Check that the join test's PartialMerge reaches the sort-merge join without an exchange, and keep AQE on by disabling the suite's runtime broadcast threshold. Name the DataFusion 56 stream as it is on main. * address review: keep Single for a PartialMerge over fully sorted input * address review: shorten the PartialMerge planning comment (cherry picked from commit d8ead26b1e7621e56c106ace11119fb30dce31e5) Adapted for branch-1.2: - Dropped the comment on `limited_to_its_share` in `fair_pool.rs`, the comment on `ORDERED_FINAL_AGGREGATE` in `spill_replay.rs`, and the paragraph in `memory_management.md`. They describe the OrderedFinalAggregateStream exemption from the fair limit, which came in with #7096 and is not on this branch. - The sort-merge join test uses a 32 MiB pool instead of 16 MiB. Without #7096 a consumer's fair limit is the pool divided by the number of consumers, and at 16 MiB the sort below the join fails before the PartialMerge runs. --- .../execution/memory_pools/spill_replay.rs | 8 +- native/core/src/execution/merge_as_partial.rs | 15 +- native/core/src/execution/planner.rs | 311 ++++++++++++++---- .../core/src/execution/planner/aggregate.rs | 59 ++++ .../comet/exec/CometAggregateSuite.scala | 98 +++++- 5 files changed, 406 insertions(+), 85 deletions(-) diff --git a/native/core/src/execution/memory_pools/spill_replay.rs b/native/core/src/execution/memory_pools/spill_replay.rs index 086e8d87374..5abf0a3a987 100644 --- a/native/core/src/execution/memory_pools/spill_replay.rs +++ b/native/core/src/execution/memory_pools/spill_replay.rs @@ -21,10 +21,10 @@ //! replays them through an `OrderedFinalAggregateStream` that has no way to spill, so a refused //! memory request there fails the task. `FinalHashAggregateStream` does this, and so does //! `OrderedFinalAggregateStream` itself, which DataFusion uses when the input is sorted on some of -//! the grouping keys. So does `SingleHashAggregateStream`, which Comet runs for a Spark -//! PartialMerge (#6857). The merge reserves read buffers for as many spill files as fit, and those -//! buffers belong to the same consumer, so the replay often finds the consumer's share already -//! taken. +//! the grouping keys, and which Comet also runs for a Spark PartialMerge over such input (#6880). +//! So does `SingleHashAggregateStream`, which Comet runs for other Spark PartialMerges (#6857). +//! The merge reserves read buffers for as many spill files as fit, and those buffers belong to the +//! same consumer, so the replay often finds the consumer's share already taken. //! //! [`SpillReplayPool`] wraps a Comet pool. When the pool refuses a request from the replay, it //! records the request with the pool's `grow` instead, which ignores `CometFairMemoryPool`'s limits diff --git a/native/core/src/execution/merge_as_partial.rs b/native/core/src/execution/merge_as_partial.rs index 4dc4c4809f7..59d24b152e4 100644 --- a/native/core/src/execution/merge_as_partial.rs +++ b/native/core/src/execution/merge_as_partial.rs @@ -27,6 +27,9 @@ //! evaluates to its state, packed into a struct. Under DataFusion's `Partial` mode, which //! outputs state, that gives merge semantics with state output. Under `Single` mode, which //! spills rather than emit a group twice, the planner unpacks the struct into the state columns. +//! It does the same under `Final` mode, which DataFusion 55 runs with an ordered aggregate that +//! spills when the input is sorted on some of the grouping keys. There, `Single` falls back to a +//! stream whose spill can ask for more than its share of memory (#6880). use std::fmt::Debug; use std::hash::{Hash, Hasher}; @@ -51,12 +54,12 @@ use datafusion_comet_spark_expr::GetStructField; use crate::execution::spark_aggregate_state::PartialMergeStateDecoder; -/// An AggregateUDF wrapper that merges state and outputs state in Partial or Single mode. +/// An AggregateUDF wrapper that merges state and outputs state in Partial, Single or Final mode. /// /// When DataFusion runs an AggregateExec in Partial mode, it calls `update_batch` /// on each accumulator and outputs `state()`. This wrapper intercepts `update_batch` /// and redirects it to `merge_batch` on the inner accumulator, effectively -/// implementing PartialMerge: merge inputs, output state. In Single mode, DataFusion +/// implementing PartialMerge: merge inputs, output state. In Single and Final modes, DataFusion /// outputs `evaluate()` instead, which returns the inner accumulator's state as a struct /// of the state fields. /// @@ -165,7 +168,7 @@ impl AggregateUDFImpl for MergeAsPartialUDF { fn return_type(&self, _arg_types: &[DataType]) -> Result { // In Partial mode, return_type isn't used for output schema (state_fields is). - // In Single mode, the wrapper evaluates to its state. + // In Single and Final modes, the wrapper evaluates to its state. Ok(DataType::Struct(self.cached_state_fields.clone())) } @@ -278,9 +281,9 @@ impl InnerAccumulatorFactory { } } -/// Projects the output of a Single-mode `aggregate` whose aggregates are all [`MergeAsPartialUDF`]s -/// into the output of a Partial-mode one: its `num_grouping` grouping columns, then one column per -/// state field, unpacked from the struct each aggregate evaluates to. +/// Projects the output of a Single- or Final-mode `aggregate` whose aggregates are all +/// [`MergeAsPartialUDF`]s into the output of a Partial-mode one: its `num_grouping` grouping +/// columns, then one column per state field, unpacked from the struct each aggregate evaluates to. pub(crate) fn unpack_state_structs( aggregate: Arc, num_grouping: usize, diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 8c240bc5dc2..268f35570a6 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -4487,7 +4487,6 @@ mod tests { #[tokio::test] async fn keyless_aggregate_without_functions_emits_one_empty_row() { use arrow::array::RecordBatchOptions; - use datafusion::common::tree_node::{Transformed, TreeNode}; let scan = Operator { op_struct: Some(OpStruct::Scan(spark_operator::Scan { @@ -4533,22 +4532,7 @@ mod tests { assert_eq!("ScanExec", planned.children[0].native_plan.name()); for batches in [vec![rows(3), rows(2)], vec![]] { - let input: Arc = - MemorySourceConfig::try_new_exec(&[batches], Arc::clone(&schema), None) - .unwrap(); - let native_plan = Arc::clone(&planned.native_plan) - .transform_up(|node| { - Ok(if node.name() == "ScanExec" { - Transformed::yes(Arc::clone(&input)) - } else { - Transformed::no(node) - }) - }) - .unwrap() - .data; - let results = collect(native_plan.execute(0, Arc::clone(&task_ctx)).unwrap()) - .await - .unwrap(); + let results = collect_over_batches(&planned, batches, Arc::clone(&task_ctx)).await; assert_eq!(1, results.iter().map(RecordBatch::num_rows).sum::()); } } @@ -4557,7 +4541,6 @@ mod tests { #[tokio::test] async fn sort_aggregate_output_is_ordered_by_grouping_keys() { use arrow::datatypes::Int32Type; - use datafusion::common::tree_node::{Transformed, TreeNode}; use spark_expression::data_type::{data_type_info::DatatypeStruct, DataTypeInfo, ListInfo}; let array_type = spark_expression::DataType { @@ -4637,21 +4620,7 @@ mod tests { ])) as ArrayRef; let batch = RecordBatch::try_new(planned.children[0].schema(), vec![Arc::clone(&keys)]).unwrap(); - let input: Arc = - MemorySourceConfig::try_new_exec(&[vec![batch.clone()]], batch.schema(), None).unwrap(); - let native_plan = Arc::clone(&planned.native_plan) - .transform_up(|node| { - Ok(if node.name() == "ScanExec" { - Transformed::yes(Arc::clone(&input)) - } else { - Transformed::no(node) - }) - }) - .unwrap() - .data; - let results = collect(native_plan.execute(0, task_ctx).unwrap()) - .await - .unwrap(); + let results = collect_over_batches(&planned, vec![batch], task_ctx).await; let output = arrow::compute::concat_batches(&results[0].schema(), &results).unwrap(); assert_eq!(keys.as_ref(), output.column(0).as_ref()); } @@ -4711,27 +4680,17 @@ mod tests { async fn partial_merge_runs_as_single_and_emits_merged_states() { use arrow::array::{BooleanArray, Int64Array}; use datafusion::assert_batches_sorted_eq; - use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode}; use spark_expression::agg_expr::ExprStruct as AggExprStruct; // A Spark PartialMerge of max(v), count(v) and first(v) grouped by k, over k and their // states. DataFusion's first has no GroupsAccumulator without an ordering. - let int64 = spark_expression::DataType { - type_id: 4, - type_info: None, - }; + let int64 = create_int64_datatype(); let boolean = spark_expression::DataType { type_id: 0, type_info: None, }; - let v = Expr { - expr_struct: Some(Unbound(spark_expression::UnboundReference { - name: "v".to_string(), - datatype: Some(int64.clone()), - })), - ..Default::default() - }; + let v = create_unbound_int64("v"); let partial_merge = spark_operator::AggregateMode::PartialMerge as i32; let aggregate = |mode: i32| Operator { plan_id: 1, @@ -4757,19 +4716,8 @@ mod tests { ..Default::default() }], agg_exprs: vec![ - spark_expression::AggExpr { - expr_struct: Some(AggExprStruct::Max(spark_expression::Max { - child: Some(v.clone()), - datatype: Some(int64.clone()), - })), - ..Default::default() - }, - spark_expression::AggExpr { - expr_struct: Some(AggExprStruct::Count(spark_expression::Count { - children: vec![v.clone()], - })), - ..Default::default() - }, + create_max_agg(&v), + create_count_agg(&v), spark_expression::AggExpr { expr_struct: Some(AggExprStruct::First(spark_expression::First { child: Some(v.clone()), @@ -4834,21 +4782,7 @@ mod tests { ], ) .unwrap(); - let input: Arc = - MemorySourceConfig::try_new_exec(&[vec![batch.clone()]], batch.schema(), None).unwrap(); - let native_plan = Arc::clone(&planned.native_plan) - .transform_up(|node| { - Ok(if node.name() == "ScanExec" { - Transformed::yes(Arc::clone(&input)) - } else { - Transformed::no(node) - }) - }) - .unwrap() - .data; - let results = collect(native_plan.execute(0, task_ctx).unwrap()) - .await - .unwrap(); + let results = collect_over_batches(&planned, vec![batch], task_ctx).await; assert_batches_sorted_eq!( [ "+-------+------------+--------------+--------------------+---------------------------+", @@ -4863,6 +4797,169 @@ mod tests { ); } + #[tokio::test] + async fn partial_merge_over_sorted_input_runs_as_final() { + use arrow::array::Int64Array; + use datafusion::assert_batches_sorted_eq; + use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode}; + use datafusion::physical_plan::InputOrderMode; + use spark_expression::agg_expr::ExprStruct as AggExprStruct; + + // A Spark PartialMerge grouped by (g, k) over a native sort on g, which DataFusion 55's + // Single would run with a GroupedHashAggregateStream that can fail to spill (#6880). + let int64 = create_int64_datatype(); + let v = create_unbound_int64("v"); + let key = |index: i32| Expr { + expr_struct: Some(Bound(spark_expression::BoundReference { + index, + datatype: Some(int64.clone()), + })), + ..Default::default() + }; + let partial_merge = spark_operator::AggregateMode::PartialMerge as i32; + let aggregate = |sort_keys: &[i32], + state_fields: Vec, + agg_exprs: Vec| { + Operator { + plan_id: 1, + children: vec![Operator { + children: vec![Operator { + op_struct: Some(OpStruct::Scan(spark_operator::Scan { + fields: [vec![int64.clone(), int64.clone()], state_fields].concat(), + source: String::new(), + })), + ..Default::default() + }], + op_struct: Some(OpStruct::Sort(spark_operator::Sort { + sort_orders: sort_keys + .iter() + .map(|&index| create_sort_order(index, 4, false, true)) + .collect(), + ..Default::default() + })), + ..Default::default() + }], + op_struct: Some(OpStruct::HashAgg(spark_operator::HashAggregate { + grouping_exprs: vec![key(0), key(1)], + expr_modes: vec![partial_merge; agg_exprs.len()], + agg_exprs, + mode: partial_merge, + initial_input_buffer_offset: 2, + ..Default::default() + })), + ..Default::default() + } + }; + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let planner = PhysicalPlanner::new(Arc::new(ctx), 0); + + // Over native max and count states, DataFusion's Final runs an ordered aggregate that + // spills, and the projection above it unpacks the states as it does above Single. + let max_count = |sort_keys: &[i32]| { + aggregate( + sort_keys, + vec![int64.clone(), int64.clone()], + vec![create_max_agg(&v), create_count_agg(&v)], + ) + }; + let (_, _, planned) = planner + .create_plan(&max_count(&[0]), &mut vec![], 1) + .unwrap(); + assert_eq!("ProjectionExec", planned.native_plan.name()); + let final_agg = planned.additional_native_plans[0] + .downcast_ref::() + .unwrap(); + assert_eq!(&AggregateMode::Final, final_agg.mode()); + assert_eq!( + &InputOrderMode::PartiallySorted(vec![0]), + final_agg.input_order_mode() + ); + + let input_schema = planned.children[0].schema(); + let batch = RecordBatch::try_new( + Arc::clone(&input_schema), + vec![ + Arc::new(Int64Array::from(vec![2, 1, 1, 2, 1])), + Arc::new(Int64Array::from(vec![1, 2, 1, 1, 2])), + Arc::new(Int64Array::from(vec![ + Some(4), + Some(10), + None, + Some(9), + Some(3), + ])), + Arc::new(Int64Array::from(vec![1, 2, 0, 3, 4])), + ], + ) + .unwrap(); + let results = collect_over_batches(&planned, vec![batch], Arc::clone(&task_ctx)).await; + assert_batches_sorted_eq!( + [ + "+-------+-------+------------+--------------+", + "| col_0 | col_1 | max[value] | count[count] |", + "+-------+-------+------------+--------------+", + "| 1 | 1 | | 0 |", + "| 1 | 2 | 10 | 6 |", + "| 2 | 1 | 9 | 4 |", + "+-------+-------+------------+--------------+", + ], + &results + ); + + // Final reads the keys from the leading columns and each state from the columns after + // them, so a PartialMerge whose input is laid out any other way stays Single. So does one + // over input sorted on every key, where Final has no spill path to offer. + let with_hash_agg = + |mut op: Operator, edit: &dyn Fn(&mut spark_operator::HashAggregate)| { + if let Some(OpStruct::HashAgg(agg)) = op.op_struct.as_mut() { + edit(agg); + } + op + }; + // Spark's partial collect_list hands over its buffers as binary, which Final would spill + // as list states into binary columns. + let collect_list = aggregate( + &[0], + vec![spark_expression::DataType { + type_id: 8, + type_info: None, + }], + vec![spark_expression::AggExpr { + expr_struct: Some(AggExprStruct::CollectList(spark_expression::CollectList { + child: Some(v.clone()), + datatype: Some(int64.clone()), + })), + ..Default::default() + }], + ); + let swapped_keys = with_hash_agg(max_count(&[0]), &|agg| { + agg.grouping_exprs = vec![key(1), key(0)]; + }); + let states_after_one_key = with_hash_agg(max_count(&[0]), &|agg| { + agg.initial_input_buffer_offset = 1; + }); + let extra_column = aggregate( + &[0], + vec![int64.clone(), int64.clone(), int64.clone()], + vec![create_max_agg(&v), create_count_agg(&v)], + ); + for (case, op) in [ + ("collect_list buffers", collect_list), + ("swapped keys", swapped_keys), + ("states after one key", states_after_one_key), + ("a column after the states", extra_column), + ("sorted on every key", max_count(&[0, 1])), + ] { + let (_, _, planned) = planner.create_plan(&op, &mut vec![], 1).unwrap(); + let single = planned.additional_native_plans[0] + .downcast_ref::() + .unwrap(); + assert_eq!(&AggregateMode::Single, single.mode(), "{case}"); + assert_ne!(&InputOrderMode::Linear, single.input_order_mode(), "{case}"); + } + } + #[tokio::test] async fn projection_prunes_filter_output_and_preserves_metrics() { use crate::execution::metrics::utils::to_native_metric_node; @@ -5022,6 +5119,74 @@ mod tests { } } + fn create_int64_datatype() -> spark_expression::DataType { + spark_expression::DataType { + type_id: 4, + type_info: None, + } + } + + /// An unbound reference to the Int64 column `name`, as a Spark aggregate function's child. + fn create_unbound_int64(name: &str) -> Expr { + Expr { + expr_struct: Some(Unbound(spark_expression::UnboundReference { + name: name.to_string(), + datatype: Some(create_int64_datatype()), + })), + ..Default::default() + } + } + + fn create_max_agg(child: &Expr) -> spark_expression::AggExpr { + spark_expression::AggExpr { + expr_struct: Some(spark_expression::agg_expr::ExprStruct::Max( + spark_expression::Max { + child: Some(child.clone()), + datatype: Some(create_int64_datatype()), + }, + )), + ..Default::default() + } + } + + fn create_count_agg(child: &Expr) -> spark_expression::AggExpr { + spark_expression::AggExpr { + expr_struct: Some(spark_expression::agg_expr::ExprStruct::Count( + spark_expression::Count { + children: vec![child.clone()], + }, + )), + ..Default::default() + } + } + + /// Runs `planned` with `batches` standing in for its single `ScanExec` and returns the + /// output batches. + async fn collect_over_batches( + planned: &crate::execution::spark_plan::SparkPlan, + batches: Vec, + task_ctx: Arc, + ) -> Vec { + use datafusion::common::tree_node::{Transformed, TreeNode}; + + let input: Arc = + MemorySourceConfig::try_new_exec(&[batches], planned.children[0].schema(), None) + .unwrap(); + let native_plan = Arc::clone(&planned.native_plan) + .transform_up(|node| { + Ok(if node.name() == "ScanExec" { + Transformed::yes(Arc::clone(&input)) + } else { + Transformed::no(node) + }) + }) + .unwrap() + .data; + collect(native_plan.execute(0, task_ctx).unwrap()) + .await + .unwrap() + } + #[tokio::test] async fn explode_evaluates_array_once_per_batch() { use arrow::datatypes::Int32Type; diff --git a/native/core/src/execution/planner/aggregate.rs b/native/core/src/execution/planner/aggregate.rs index d0c26bddd9e..19b4271ace1 100644 --- a/native/core/src/execution/planner/aggregate.rs +++ b/native/core/src/execution/planner/aggregate.rs @@ -247,6 +247,34 @@ impl OperatorBuilder for HashAggBuilder { }) .collect(); + // DataFusion 55 runs Single over input sorted on some of its keys with the legacy + // GroupedHashAggregateStream, whose spill can fail against the fair limit (#6880). Final + // runs an OrderedFinalAggregateStream there, which spills, and which the fair pool exempts + // from that limit (`limited_to_its_share`). MergeAsPartial evaluates to its state under + // Final as under Single. + // Only partially sorted input switches: with fully sorted input Final has no spill path + // and would halve the limit of a sort below it. Final also reads states from the columns + // after the keys, so the input must be the keys followed by native states (Spark's binary + // collect buffers keep Single, and arrive unordered anyway). + // Drop this on DataFusion 56, where the stream is ClusteredFinalAggregateStream + // (apache/datafusion#24697), and update the name prefix in `spill_replay.rs` (#6410). + let mode = if proto_mode == ProtoAggregateMode::PartialMerge + && input_is_keys_then_states(agg, &aggr_expr, &schema)? + && { + let keys = group_by.input_exprs(); + let sorted_on = child + .native_plan + .equivalence_properties() + .find_longest_permutation(&keys)? + .1; + // DataFusion's InputOrderMode is PartiallySorted when some, not all, keys match. + !sorted_on.is_empty() && sorted_on.len() < keys.len() + } { + DFAggregateMode::Final + } else { + mode + }; + let aggregate: Arc = Arc::new( datafusion::physical_plan::aggregates::AggregateExec::try_new( mode, @@ -309,6 +337,37 @@ impl OperatorBuilder for HashAggBuilder { } } +/// Whether a PartialMerge's input is its grouping keys, in order, followed by the states of +/// its MergeAsPartial aggregates in their native types, which DataFusion's Final assumes. A +/// Spark partial aggregate's collect buffers, for example, arrive as binary. +fn input_is_keys_then_states( + agg: &datafusion_comet_proto::spark_operator::HashAggregate, + aggr_expr: &[Arc], + input_schema: &SchemaRef, +) -> Result { + let num_keys = agg.grouping_exprs.len(); + let keys_lead = agg.initial_input_buffer_offset as usize == num_keys + && agg.grouping_exprs.iter().enumerate().all(|(idx, expr)| { + matches!(&expr.expr_struct, Some(ExprStruct::Bound(bound)) if bound.index as usize == idx) + }); + if !keys_lead { + return Ok(false); + } + let mut input_types = input_schema + .fields() + .iter() + .skip(num_keys) + .map(|f| f.data_type()); + for expr in aggr_expr { + for field in expr.state_fields()? { + if input_types.next() != Some(field.data_type()) { + return Ok(false); + } + } + } + Ok(input_types.next().is_none()) +} + impl PhysicalPlanner { /// Create a DataFusion physical aggregate expression from Spark physical aggregate expression fn create_agg_expr( 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 47eb0203d9e..359b2cafe68 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -31,9 +31,9 @@ import org.apache.spark.sql.catalyst.expressions.{Cast, Literal} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Corr, Final, Partial, PartialMerge, RegrR2} import org.apache.spark.sql.catalyst.optimizer.EliminateSorts import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning} -import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec, CometSortAggregateExec, CometSortExec} +import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometFilterExec, CometHashAggregateExec, CometNativeExec, CometProjectExec, CometSortAggregateExec, CometSortExec, CometSortMergeJoinExec} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} -import org.apache.spark.sql.execution.LocalTableScanExec +import org.apache.spark.sql.execution.{LocalTableScanExec, SparkPlan} import org.apache.spark.sql.execution.SQLExecution import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, ShuffleQueryStageExec} import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} @@ -301,6 +301,100 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + // Writes 125,000 rows over four values of g, with 62,500 distinct k that each fall in one g, + // then runs `f` in each of the pools. The PartialMerge keyed by (g, k) has to spill several runs + // of its 62,500 groups, and its input is sorted on g, a prefix of its keys, so DataFusion runs + // it as an ordered final aggregate (issue #6880). Each g has 15,625 distinct k and the largest v + // with id % 4 = g, so a key emitted twice across a change of g shows up in the answer. + private def withPartialMergeOverSortedPrefix(poolsMiB: Seq[Int])(f: String => Unit): Unit = { + withTempPath { dir => + val path = dir.getCanonicalPath + spark + .range(0, 125000, 1, 1) + .selectExpr("id % 4 AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k') AS k", "id AS v") + .write + .parquet(path + "/data") + spark.range(0, 4).selectExpr("id AS g").write.parquet(path + "/keys") + for (poolMiB <- poolsMiB) { + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true", + "spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> (poolMiB / 2048.0).toString) { + withClue(s"With a $poolMiB MiB pool: ") { + f(path) + } + } + } + } + } + + private val partialMergeOverSortedPrefixAnswer = + (0 until 4).map(g => Row(g.toLong, 124996L + g, 15625L)) + + test("PartialMerge over input sorted on a prefix of its keys spills (issue #6880)") { + withPartialMergeOverSortedPrefix(Seq(2, 8)) { path => + // The repartition on g satisfies the PartialMerge keyed by (g, k), so the sort, the partial + // aggregate and the PartialMerge run in one native plan with no exchange between them. The + // optimizer drops the sort under the aggregate unless EliminateSorts is excluded. + withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> EliminateSorts.ruleName) { + val df = spark.read + .parquet(path + "/data") + .repartition(1, col("g")) + .sortWithinPartitions("g") + .groupBy("g") + .agg(expr("max(v)"), count_distinct(col("k"))) + .orderBy("g") + checkCometAnswer(df, partialMergeOverSortedPrefixAnswer) + + assert( + spilledAggregates(df, PartialMerge).exists(_.child match { + case partial: CometHashAggregateExec => partial.child.isInstanceOf[CometSortExec] + case _ => false + }), + "Expected a native PartialMerge over a partial aggregate over a sort:\n" + + df.queryExecution.executedPlan) + } + } + } + + test("PartialMerge over a sort-merge join on a prefix of its keys spills (issue #6880)") { + // The sorts under the join need more than a 2 or 8 MiB pool's fair share, so this uses 32 MiB. + withPartialMergeOverSortedPrefix(Seq(32)) { path => + // With the default rules, the join's hash partitioning on g satisfies the PartialMerge keyed + // by (g, k), so the join, the partial aggregate and the PartialMerge run in one native plan, + // and the PartialMerge sees its input sorted on g. The suite lets AQE broadcast up to 1g at + // runtime, which would turn the join into a broadcast hash join, because the keys are tiny. + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + val df = spark.read + .parquet(path + "/data") + .join(spark.read.parquet(path + "/keys"), "g") + .groupBy("g") + .agg(expr("max(v)"), count_distinct(col("k"))) + .orderBy("g") + checkCometAnswer(df, partialMergeOverSortedPrefixAnswer) + + // An exchange between the join and the partial aggregate would leave the PartialMerge + // over unsorted input, where it would still spill. + def reachesJoinWithoutExchange(plan: SparkPlan): Boolean = plan match { + case _: CometSortMergeJoinExec => true + case _: CometShuffleExchangeExec => false + case other => other.children.exists(reachesJoinWithoutExchange) + } + assert( + spilledAggregates(df, PartialMerge).exists(_.child match { + case partial: CometHashAggregateExec => reachesJoinWithoutExchange(partial.child) + case _ => false + }), + "Expected a native PartialMerge over a partial aggregate over a sort-merge join:\n" + + df.queryExecution.executedPlan) + } + } + } + test("collect_list PartialMerge short of memory keeps one row per key (issue #6857)") { // The collect_list states fill the PartialMerge's memory too. Over a Spark partial, the // PartialMerge decodes Spark's buffers as it merges them, but reads back the native states