From 2f35370585df21dea3d9dcc4d667c23909073dc5 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 11 Oct 2026 09:35:07 -0600 Subject: [PATCH 1/2] fix: keep one row per key in a native PartialMerge that runs short of memory Comet ran Spark's PartialMerge aggregate as DataFusion Partial, which emits the groups it holds when it runs short of memory and starts again. In Spark's single-distinct plan the stage above the PartialMerge counts its rows, so count(DISTINCT) counted every copy of a key. Run a PartialMerge as DataFusion Single instead, which spills. MergeAsPartial evaluates to a struct of its state, and a projection unpacks it into the state columns that the Partial plan emitted. SpillReplayPool also covers SingleHashAggregateStream, whose spill replay starves the same way as the final aggregates' do (#6254). Closes #6857. --- native/core/src/execution/jni_api.rs | 2 +- .../execution/memory_pools/spill_replay.rs | 19 +- native/core/src/execution/merge_as_partial.rs | 113 ++++++++--- native/core/src/execution/planner.rs | 183 ++++++++++++++++-- .../comet/exec/CometAggregateSuite.scala | 118 +++++++++-- 5 files changed, 370 insertions(+), 65 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 0812b9f7087..a2646dd8daf 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -963,7 +963,7 @@ fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operato if count.children.len() == 1) }) } - // PartialMerge is represented as native Partial, but consumes states, not rows. + // PartialMerge consumes states, not rows. _ => false, }, _ => true, diff --git a/native/core/src/execution/memory_pools/spill_replay.rs b/native/core/src/execution/memory_pools/spill_replay.rs index e7d947e38c1..086e8d87374 100644 --- a/native/core/src/execution/memory_pools/spill_replay.rs +++ b/native/core/src/execution/memory_pools/spill_replay.rs @@ -21,7 +21,8 @@ //! 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. The merge reserves read buffers for as many spill files as fit, and those +//! 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. //! @@ -31,9 +32,9 @@ //! //! This relies on the following in DataFusion 55.1, which a DataFusion upgrade has to re-check: //! -//! - The two aggregates name their consumers `FinalHashAggregateStream[{partition}]` and -//! `OrderedFinalAggregateStream[{partition}]`, and their merge and replay reserve through sibling -//! reservations of that consumer. +//! - The three aggregates name their consumers `FinalHashAggregateStream[{partition}]`, +//! `OrderedFinalAggregateStream[{partition}]` and `SingleHashAggregateStream[{partition}]`, and +//! their merge and replay reserve through sibling reservations of that consumer. //! - A final aggregate grows one reservation while another of its reservations holds memory only //! during the replay, while the merge holds its read buffers. Before that, the aggregate's table //! is its only reservation holding memory, so a refusal stands and makes it spill. The merge @@ -45,7 +46,8 @@ //! //! Remove this, as #6583 describes, once Comet's DataFusion has apache/datafusion#25383, which //! leaves the replay room when the merge picks its files. The #6254 tests in `CometAggregateSuite` -//! fail without this on DataFusion 55.1, so they show whether the replay still needs it. +//! fail without this on DataFusion 55.1, and its PartialMerge read-back test for #6857 fails +//! without it on 55.2, so they show whether the replay still needs it. use std::collections::HashMap; use std::fmt; @@ -174,11 +176,12 @@ pub(super) fn unwrap_spill_replay(pool: &Arc) -> Option<&Arc bool { let name = consumer.name(); name.starts_with("FinalHashAggregateStream[") || name.starts_with("OrderedFinalAggregateStream[") + || name.starts_with("SingleHashAggregateStream[") } #[cfg(test)] @@ -353,14 +356,14 @@ mod tests { for name in [ "FinalHashAggregateStream[3]", "OrderedFinalAggregateStream[0]", + "SingleHashAggregateStream[0]", ] { assert!(is_final_aggregate(&MemoryConsumer::new(name)), "{name}"); } - // Partial aggregates emit early instead of spilling, and Comet never plans single ones. + // Partial aggregates emit early instead of spilling. for name in [ "PartialHashAggregateStream[0]", "OrderedPartialAggregateStream[0]", - "SingleHashAggregateStream[0]", "ExternalSorterMerge[0]", ] { assert!(!is_final_aggregate(&MemoryConsumer::new(name)), "{name}"); diff --git a/native/core/src/execution/merge_as_partial.rs b/native/core/src/execution/merge_as_partial.rs index 7f483a761c5..356dfac8d2e 100644 --- a/native/core/src/execution/merge_as_partial.rs +++ b/native/core/src/execution/merge_as_partial.rs @@ -20,18 +20,21 @@ //! Spark's PartialMerge mode merges intermediate state buffers and outputs intermediate //! state (not final values). DataFusion has no equivalent mode — `Partial` calls //! `update_batch` and outputs state, while `Final` calls `merge_batch` and outputs -//! evaluated results. +//! evaluated results. `PartialReduce` merges state and outputs state, but from DataFusion 56 +//! it emits a group more than once when it runs short of memory, as `Partial` does. //! -//! This wrapper bridges the gap: it operates under DataFusion's `Partial` mode (which -//! outputs state) but redirects `update_batch` calls to `merge_batch`, giving merge -//! semantics with state output. +//! This wrapper bridges the gap: it redirects `update_batch` calls to `merge_batch`, and it +//! 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. use std::fmt::Debug; use std::hash::{Hash, Hasher}; +use std::sync::Arc; -use arrow::array::{ArrayRef, BooleanArray}; -use arrow::datatypes::{DataType, Field, FieldRef}; -use datafusion::common::Result; +use arrow::array::{ArrayRef, BooleanArray, StructArray}; +use arrow::datatypes::{DataType, Field, FieldRef, Fields}; +use datafusion::common::{internal_err, Result}; use datafusion::logical_expr::function::AccumulatorArgs; use datafusion::logical_expr::function::StateFieldsArgs; use datafusion::logical_expr::{ @@ -39,16 +42,21 @@ use datafusion::logical_expr::{ Signature, Volatility, }; use datafusion::physical_expr::aggregate::AggregateFunctionExpr; +use datafusion::physical_expr::{expressions::Column, PhysicalExpr}; +use datafusion::physical_plan::{projection::ProjectionExec, ExecutionPlan}; use datafusion::scalar::ScalarValue; +use datafusion_comet_spark_expr::GetStructField; use crate::execution::spark_aggregate_state::PartialMergeStateDecoder; -/// An AggregateUDF wrapper that gives merge semantics in Partial mode. +/// An AggregateUDF wrapper that merges state and outputs state in Partial or Single 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. +/// implementing PartialMerge: merge inputs, output state. In Single mode, DataFusion +/// outputs `evaluate()` instead, which returns the inner accumulator's state as a struct +/// of the state fields. /// /// We store the inner AggregateUDF (not the AggregateFunctionExpr) to avoid keeping /// references to UnboundColumn expressions that would panic if evaluated. @@ -56,10 +64,12 @@ use crate::execution::spark_aggregate_state::PartialMergeStateDecoder; pub struct MergeAsPartialUDF { /// The inner aggregate UDF, cloned from the original expression. inner_udf: AggregateUDF, - /// Pre-computed return type from the original expression. - return_type: DataType, - /// Pre-computed state fields from the original expression. - cached_state_fields: Vec, + /// The original expression's return field, which the inner accumulator is built with, since + /// the wrapper returns its state instead. + inner_return_field: FieldRef, + /// Pre-computed state fields from the original expression, which are also the fields of + /// the struct the wrapper evaluates to. + cached_state_fields: Fields, /// Cached signature that accepts state field types. signature: Signature, /// Override argument fields used only when constructing the inner accumulator. @@ -87,8 +97,7 @@ impl Hash for MergeAsPartialUDF { impl MergeAsPartialUDF { pub fn new(inner_expr: &AggregateFunctionExpr) -> Result { let name = format!("merge_as_partial_{}", inner_expr.name()); - let return_type = inner_expr.field().data_type().clone(); - let cached_state_fields = inner_expr.state_fields()?; + let cached_state_fields = Fields::from(inner_expr.state_fields()?); let accumulator_expr_fields = Self::accumulator_expr_fields(inner_expr, &cached_state_fields); let state_decoder = PartialMergeStateDecoder::try_new(inner_expr, &cached_state_fields)?; @@ -99,7 +108,7 @@ impl MergeAsPartialUDF { Ok(Self { inner_udf: inner_expr.fun().clone(), - return_type, + inner_return_field: inner_expr.field(), cached_state_fields, signature, accumulator_expr_fields, @@ -136,6 +145,7 @@ impl MergeAsPartialUDF { .as_deref() .unwrap_or(args.expr_fields); AccumulatorArgs { + return_field: Arc::clone(&self.inner_return_field), expr_fields, ..args } @@ -153,14 +163,20 @@ 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). - // Return the inner function's return type for consistency. - Ok(self.return_type.clone()) + // In Single mode, the wrapper evaluates to its state. + Ok(DataType::Struct(self.cached_state_fields.clone())) + } + + fn is_nullable(&self) -> bool { + // The state struct itself is never null, so each state field unpacked from it keeps its + // own nullability. + false } fn state_fields(&self, _args: StateFieldsArgs) -> Result> { // Cached at construction: state schema depends on the inner aggregate's // return type, not on StateFieldsArgs. - Ok(self.cached_state_fields.clone()) + Ok(self.cached_state_fields.to_vec()) } fn accumulator(&self, args: AccumulatorArgs) -> Result> { @@ -170,6 +186,7 @@ impl AggregateUDFImpl for MergeAsPartialUDF { Ok(Box::new(MergeAsPartialAccumulator { inner: inner_acc, state_decoder: self.state_decoder.clone(), + state_fields: self.cached_state_fields.clone(), })) } @@ -188,6 +205,7 @@ impl AggregateUDFImpl for MergeAsPartialUDF { Ok(Box::new(MergeAsPartialGroupsAccumulator { inner: inner_acc, state_decoder: self.state_decoder.clone(), + state_fields: self.cached_state_fields.clone(), })) } @@ -204,10 +222,42 @@ impl AggregateUDFImpl for MergeAsPartialUDF { } } -/// Accumulator wrapper that redirects update_batch to merge_batch. +/// 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. +pub(crate) fn unpack_state_structs( + aggregate: Arc, + num_grouping: usize, +) -> Result> { + let schema = aggregate.schema(); + let mut exprs: Vec<(Arc, String)> = Vec::new(); + for (idx, field) in schema.fields().iter().enumerate() { + let column: Arc = Arc::new(Column::new(field.name(), idx)); + if idx < num_grouping { + exprs.push((column, field.name().clone())); + continue; + } + let DataType::Struct(state_fields) = field.data_type() else { + return internal_err!( + "Expected a struct of aggregate state in column {idx}, got {}", + field.data_type() + ); + }; + for (ordinal, state_field) in state_fields.iter().enumerate() { + exprs.push(( + Arc::new(GetStructField::new(Arc::clone(&column), ordinal)), + state_field.name().clone(), + )); + } + } + Ok(Arc::new(ProjectionExec::try_new(exprs, aggregate)?)) +} + +/// Accumulator wrapper that redirects update_batch to merge_batch and evaluates to its state. struct MergeAsPartialAccumulator { inner: Box, state_decoder: PartialMergeStateDecoder, + state_fields: Fields, } impl Debug for MergeAsPartialAccumulator { @@ -229,7 +279,17 @@ impl Accumulator for MergeAsPartialAccumulator { } fn evaluate(&mut self) -> Result { - self.inner.evaluate() + let state = self + .inner + .state()? + .iter() + .map(ScalarValue::to_array) + .collect::>>()?; + Ok(ScalarValue::Struct(Arc::new(StructArray::try_new( + self.state_fields.clone(), + state, + None, + )?))) } fn state(&mut self) -> Result> { @@ -241,10 +301,12 @@ impl Accumulator for MergeAsPartialAccumulator { } } -/// GroupsAccumulator wrapper that redirects update_batch to merge_batch. +/// GroupsAccumulator wrapper that redirects update_batch to merge_batch and evaluates to its +/// state. struct MergeAsPartialGroupsAccumulator { inner: Box, state_decoder: PartialMergeStateDecoder, + state_fields: Fields, } impl Debug for MergeAsPartialGroupsAccumulator { @@ -290,7 +352,12 @@ impl GroupsAccumulator for MergeAsPartialGroupsAccumulator { } fn evaluate(&mut self, emit_to: EmitTo) -> Result { - self.inner.evaluate(emit_to) + let state = self.inner.state(emit_to)?; + Ok(Arc::new(StructArray::try_new( + self.state_fields.clone(), + state, + None, + )?)) } fn state(&mut self, emit_to: EmitTo) -> Result> { diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index cc9157f8468..57373b4b130 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1575,13 +1575,17 @@ impl PhysicalPlanner { let mode = match proto_mode { ProtoAggregateMode::Partial => DFAggregateMode::Partial, ProtoAggregateMode::Final => DFAggregateMode::Final, - // PartialMerge: Partial + MergeAsPartial - ProtoAggregateMode::PartialMerge => DFAggregateMode::Partial, + // Spark's distinct-aggregate plan counts the rows of a PartialMerge, so it must + // emit each key once. DataFusion's Partial emits what it holds when it runs + // short of memory and starts again, while its Single spills. So run the + // PartialMerge as Single, whose MergeAsPartial aggregates merge their input and + // evaluate to their states, and unpack those states below. + ProtoAggregateMode::PartialMerge => DFAggregateMode::Single, }; // Check if any expression uses PartialMerge mode. When present, // those expressions are wrapped with MergeAsPartial to get merge - // semantics inside a Partial-mode AggregateExec. + // semantics with state output. let partial_merge_value = ProtoAggregateMode::PartialMerge as i32; let has_partial_merge = proto_mode == ProtoAggregateMode::PartialMerge || agg.expr_modes.contains(&partial_merge_value); @@ -1680,6 +1684,18 @@ impl PhysicalPlanner { )?, ); + // A PartialMerge runs as Single, so unpack the states it evaluates to into the + // state columns that Spark expects of it. + let mut output = Arc::clone(&aggregate); + let mut additional_native_plans = vec![]; + if proto_mode == ProtoAggregateMode::PartialMerge { + output = crate::execution::merge_as_partial::unpack_state_structs( + output, + agg.grouping_exprs.len(), + )?; + additional_native_plans.push(Arc::clone(&aggregate)); + } + // Spark's SortAggregateExec reports its output as ordered by the grouping keys, // and Spark may have removed a sort above it on that basis. DataFusion emits // groups in input order only when it sees the input sorted on them. When that @@ -1687,11 +1703,11 @@ impl PhysicalPlanner { // relation that reaches native code as an unordered ScanExec, the hash table can // emit groups out of order, so sort the aggregate output instead. if agg.ordered_by_grouping_keys && !agg.grouping_exprs.is_empty() { - let aggregate_schema = aggregate.schema(); + let output_schema = output.schema(); let ordering: Vec = (0..agg.grouping_exprs.len()) .map(|idx| { PhysicalSortExpr::new( - Arc::new(Column::new(aggregate_schema.field(idx).name(), idx)), + Arc::new(Column::new(output_schema.field(idx).name(), idx)), SortOptions { descending: false, nulls_first: true, @@ -1699,31 +1715,25 @@ impl PhysicalPlanner { ) }) .collect(); - if !aggregate + if !output .equivalence_properties() .ordering_satisfy(ordering.clone())? { - let sort: Arc = Arc::new(SortExec::new( - LexOrdering::new(ordering).unwrap(), - Arc::clone(&aggregate), - )); - return Ok(( - scans, - shuffle_scans, - Arc::new(SparkPlan::new_with_additional( - spark_plan.plan_id, - sort, - vec![child], - vec![aggregate], - )), - )); + additional_native_plans.push(Arc::clone(&output)); + output = + Arc::new(SortExec::new(LexOrdering::new(ordering).unwrap(), output)); } } Ok(( scans, shuffle_scans, - Arc::new(SparkPlan::new(spark_plan.plan_id, aggregate, vec![child])), + Arc::new(SparkPlan::new_with_additional( + spark_plan.plan_id, + output, + vec![child], + additional_native_plans, + )), )) } @@ -6316,6 +6326,137 @@ mod tests { assert_eq!("AggregateExec", planned.native_plan.name()); } + #[tokio::test] + async fn partial_merge_runs_as_single_and_emits_merged_states() { + use arrow::array::Int64Array; + 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) and count(v) grouped by k, over [k, max state, count state]. + let int64 = spark_expression::DataType { + type_id: 4, + type_info: None, + }; + let v = Expr { + expr_struct: Some(Unbound(spark_expression::UnboundReference { + name: "v".to_string(), + datatype: Some(int64.clone()), + })), + ..Default::default() + }; + let partial_merge = spark_operator::AggregateMode::PartialMerge as i32; + let aggregate = |mode: i32| Operator { + plan_id: 1, + children: vec![Operator { + op_struct: Some(OpStruct::Scan(spark_operator::Scan { + fields: vec![int64.clone(); 3], + source: String::new(), + })), + ..Default::default() + }], + op_struct: Some(OpStruct::HashAgg(spark_operator::HashAggregate { + grouping_exprs: vec![Expr { + expr_struct: Some(Bound(spark_expression::BoundReference { + index: 0, + datatype: Some(int64.clone()), + })), + ..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() + }, + ], + mode, + expr_modes: vec![partial_merge; 2], + initial_input_buffer_offset: 1, + ..Default::default() + })), + ..Default::default() + }; + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let planner = PhysicalPlanner::new(Arc::new(ctx), 0); + + // DataFusion's Single spills rather than emit a key twice, and the projection above it + // unpacks the states into the columns that the Partial it replaces emitted. + let (_, _, planned) = planner + .create_plan(&aggregate(partial_merge), &mut vec![], 1) + .unwrap(); + assert_eq!("ProjectionExec", planned.native_plan.name()); + assert_eq!(1, planned.additional_native_plans.len()); + let single = planned.additional_native_plans[0] + .downcast_ref::() + .unwrap(); + assert_eq!(&AggregateMode::Single, single.mode()); + let (_, _, partial) = planner + .create_plan( + &aggregate(spark_operator::AggregateMode::Partial as i32), + &mut vec![], + 1, + ) + .unwrap(); + assert_eq!("AggregateExec", partial.native_plan.name()); + assert_eq!(partial.native_plan.schema(), planned.native_plan.schema()); + + let batch = RecordBatch::try_new( + planned.children[0].schema(), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 1, 2, 3])), + Arc::new(Int64Array::from(vec![ + Some(10), + Some(5), + Some(7), + None, + Some(1), + ])), + Arc::new(Int64Array::from(vec![1, 2, 3, 0, 4])), + ], + ) + .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 output = arrow::compute::concat_batches(&results[0].schema(), &results).unwrap(); + let column = |idx: usize| { + output + .column(idx) + .as_any() + .downcast_ref::() + .unwrap() + .clone() + }; + let (keys, maxes, counts) = (column(0), column(1), column(2)); + let mut rows: Vec<_> = (0..output.num_rows()) + .map(|row| (keys.value(row), maxes.value(row), counts.value(row))) + .collect(); + rows.sort(); + assert_eq!(vec![(1, 10, 4), (2, 5, 2), (3, 1, 4)], rows); + } + #[tokio::test] async fn projection_prunes_filter_output_and_preserves_metrics() { use crate::execution::metrics::utils::to_native_metric_node; 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 a8bd439c188..f4993ff34e6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -28,7 +28,7 @@ import org.apache.spark.{CometListenerBusUtils, SparkConf} import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd} import org.apache.spark.sql.{Column, CometTestBase, DataFrame, QueryTest, Row} import org.apache.spark.sql.catalyst.expressions.Cast -import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge} 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} @@ -156,19 +156,21 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } /** - * Returns the native final aggregates in the plan that `df` last ran, after checking that there - * is one and that each spilled. + * Returns the native aggregates in the plan that `df` last ran whose aggregate functions all + * run in `mode`, after checking that there is one and that each spilled. */ - private def spilledFinalAggregates(df: DataFrame): Seq[CometHashAggregateExec] = { + private def spilledAggregates( + df: DataFrame, + mode: AggregateMode): Seq[CometHashAggregateExec] = { val plan = df.queryExecution.executedPlan - val finalAggregates = collect(plan) { - case aggregate: CometHashAggregateExec if aggregate.modes.contains(Final) => aggregate + val aggregates = collect(plan) { + case aggregate: CometHashAggregateExec if aggregate.modes == Seq(mode) => aggregate } - assert(finalAggregates.nonEmpty, s"Expected a native final aggregate:\n$plan") + assert(aggregates.nonEmpty, s"Expected a native $mode aggregate:\n$plan") assert( - finalAggregates.forall(_.metrics("spill_count").value > 0), - "The final aggregate did not spill") - finalAggregates + aggregates.forall(_.metrics("spill_count").value > 0), + s"The $mode aggregate did not spill") + aggregates } test("final aggregate that has spilled reads its spill files back (issue #6254)") { @@ -193,7 +195,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } // checkCometAnswer runs the query once, so the plan read below is the one that ran. checkCometAnswer(df, expected) - spilledFinalAggregates(df) + spilledAggregates(df, Final) } } } @@ -233,7 +235,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { // Only a sort in the same native plan keeps the final aggregate's input sorted. assert( - spilledFinalAggregates(df).exists(_.child match { + spilledAggregates(df, Final).exists(_.child match { case partial: CometHashAggregateExec => partial.child.isInstanceOf[CometSortExec] case _ => false }), @@ -243,6 +245,98 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + /** + * Runs `f` with a native memory pool small enough to make the PartialMerge stage of a distinct + * aggregate run short of memory: about 4 MiB, 0.002 of the suite's 2g off-heap. + */ + private def withDistinctMergeMemoryPressure(f: => Unit): Unit = { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002", + CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true", + "spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536")(f) + } + + test("PartialMerge that runs short of memory keeps one row per key (issue #6857)") { + // Spark plans count(DISTINCT d) next to sum(v) with a PartialMerge keyed by (g, d), and the + // stage above it counts the PartialMerge's rows, so the PartialMerge has to spill rather than + // emit a key twice. + withTempPath { dir => + val path = dir.getCanonicalPath + spark + .range(0, 400000, 1, 1) + .selectExpr("id % 100 AS g", "id % 200000 AS d", "id AS v") + .write + .parquet(path) + withDistinctMergeMemoryPressure { + val df = spark.read.parquet(path).groupBy("g").agg(sum("v"), count_distinct(col("d"))) + // Each g holds the 4000 ids g + 100 * i, whose d values pair up into 2000 distinct ones. + val expected = (0 until 100).map(g => Row(g.toLong, 4000L * g + 799800000L, 2000L)) + checkCometAnswer(df, expected) + spilledAggregates(df, PartialMerge) + } + } + } + + test("PartialMerge that has spilled reads its spill files back (issue #6857)") { + withTempPath { dir => + val path = dir.getCanonicalPath + spark + .range(0, 125000, 1, 1) + .selectExpr("0L AS g", "lpad(cast(id % 62500 AS STRING), 128, 'k') AS k", "id AS v") + .write + .parquet(path) + // A 2 MiB pool, shared with the aggregate above it, makes the PartialMerge keyed by (g, k) + // spill several runs of its 62,500 groups. As in the #6254 tests, merging the runs then + // takes nearly all of its share, and reading them back must not fail the task. + withSQLConf( + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "1024", + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> (2.0 / 2048).toString) { + val df = + spark.read.parquet(path).groupBy("g").agg(expr("max(v)"), count_distinct(col("k"))) + checkCometAnswer(df, Seq(Row(0L, 124999L, 62500L))) + spilledAggregates(df, PartialMerge) + } + } + } + + 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 + // it spilled. + def padded(i: Int): String = "x" * (128 - i.toString.length) + i + val rows = (0 until 32768).map(i => (i % 4096, i % 8192, padded(i))) + // Each k holds the ids k + 4096 * j, whose d values are k and k + 4096. + val expected = (0 until 4096).map { k => + Row(k, 2L, (0 until 8).map(j => padded(k + 4096 * j)).sorted) + } + withTempPath { dir => + val path = dir.getCanonicalPath + rows.toDF("k", "d", "s").write.parquet(path) + for (sparkPartial <- Seq(false, true)) { + withDistinctMergeMemoryPressure { + val source = if (sparkPartial) rows.toDF("k", "d", "s") else spark.read.parquet(path) + val df = + source.groupBy("k").agg(count_distinct(col("d")), sort_array(collect_list("s"))) + checkCometAnswer(df, expected) + spilledAggregates(df, PartialMerge) + val sparkPartials = collect(df.queryExecution.executedPlan) { + case agg: ObjectHashAggregateExec + if agg.aggregateExpressions.forall(_.mode == Partial) => + agg + } + assert( + sparkPartials.nonEmpty == sparkPartial, + "Expected a Spark partial aggregate only over the local rows:\n" + + df.queryExecution.executedPlan) + } + } + } + } + test("collect_list over struct with non-nullable fields") { // Building a struct from non-nullable columns yields non-nullable struct fields. Native // collect_list derives its declared output element type from the argument's declared type, From 2ce73e59e1be2c4df85275641f21ab7b90df9c57 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 11 Oct 2026 11:51:07 -0600 Subject: [PATCH 2/2] Emit PartialMerge states a column at a time for row-only aggregates MergeAsPartial now wraps an inner aggregate that has no GroupsAccumulator (first, last and any_value without an ordering) in a GroupsAccumulatorAdapter itself. Under Single mode, DataFusion otherwise wrapped the wrapper's row accumulators and evaluated every group to a one-row struct. Also write the planner's two output wrappers the same way, check the planner test's output with assert_batches_sorted_eq and cover first there, and drop two settings that the memory-pressure tests don't read. --- native/core/src/execution/merge_as_partial.rs | 70 +++++++++++++++++-- native/core/src/execution/planner.rs | 69 ++++++++++++------ .../comet/exec/CometAggregateSuite.scala | 9 ++- 3 files changed, 113 insertions(+), 35 deletions(-) diff --git a/native/core/src/execution/merge_as_partial.rs b/native/core/src/execution/merge_as_partial.rs index 356dfac8d2e..4dc4c4809f7 100644 --- a/native/core/src/execution/merge_as_partial.rs +++ b/native/core/src/execution/merge_as_partial.rs @@ -33,7 +33,7 @@ use std::hash::{Hash, Hasher}; use std::sync::Arc; use arrow::array::{ArrayRef, BooleanArray, StructArray}; -use arrow::datatypes::{DataType, Field, FieldRef, Fields}; +use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema}; use datafusion::common::{internal_err, Result}; use datafusion::logical_expr::function::AccumulatorArgs; use datafusion::logical_expr::function::StateFieldsArgs; @@ -42,7 +42,9 @@ use datafusion::logical_expr::{ Signature, Volatility, }; use datafusion::physical_expr::aggregate::AggregateFunctionExpr; -use datafusion::physical_expr::{expressions::Column, PhysicalExpr}; +use datafusion::physical_expr::{ + expressions::Column, GroupsAccumulatorAdapter, PhysicalExpr, PhysicalSortExpr, +}; use datafusion::physical_plan::{projection::ProjectionExec, ExecutionPlan}; use datafusion::scalar::ScalarValue; use datafusion_comet_spark_expr::GetStructField; @@ -190,18 +192,26 @@ impl AggregateUDFImpl for MergeAsPartialUDF { })) } - fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - self.inner_udf - .groups_accumulator_supported(self.inner_accumulator_args(args)) + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + // An inner aggregate without a GroupsAccumulator gets an adapter below, so that Single + // mode emits its states a column at a time rather than evaluating each group to a struct. + true } fn create_groups_accumulator( &self, args: AccumulatorArgs, ) -> Result> { - let inner_acc = self + let inner_args = self.inner_accumulator_args(args); + let inner_acc: Box = if self .inner_udf - .create_groups_accumulator(self.inner_accumulator_args(args))?; + .groups_accumulator_supported(inner_args.clone()) + { + self.inner_udf.create_groups_accumulator(inner_args)? + } else { + let factory = InnerAccumulatorFactory::new(&self.inner_udf, &inner_args); + Box::new(GroupsAccumulatorAdapter::new(move || factory.create())) + }; Ok(Box::new(MergeAsPartialGroupsAccumulator { inner: inner_acc, state_decoder: self.state_decoder.clone(), @@ -222,6 +232,52 @@ impl AggregateUDFImpl for MergeAsPartialUDF { } } +/// Owned copies of the arguments that build an inner row accumulator, which a +/// `GroupsAccumulatorAdapter` does for each new group, after `AccumulatorArgs` has gone. +struct InnerAccumulatorFactory { + udf: AggregateUDF, + return_field: FieldRef, + schema: Schema, + ignore_nulls: bool, + order_bys: Vec, + is_reversed: bool, + name: String, + is_distinct: bool, + exprs: Vec>, + expr_fields: Vec, +} + +impl InnerAccumulatorFactory { + fn new(udf: &AggregateUDF, args: &AccumulatorArgs) -> Self { + Self { + udf: udf.clone(), + return_field: Arc::clone(&args.return_field), + schema: args.schema.clone(), + ignore_nulls: args.ignore_nulls, + order_bys: args.order_bys.to_vec(), + is_reversed: args.is_reversed, + name: args.name.to_string(), + is_distinct: args.is_distinct, + exprs: args.exprs.to_vec(), + expr_fields: args.expr_fields.to_vec(), + } + } + + fn create(&self) -> Result> { + self.udf.accumulator(AccumulatorArgs { + return_field: Arc::clone(&self.return_field), + schema: &self.schema, + ignore_nulls: self.ignore_nulls, + order_bys: &self.order_bys, + is_reversed: self.is_reversed, + name: &self.name, + is_distinct: self.is_distinct, + exprs: &self.exprs, + expr_fields: &self.expr_fields, + }) + } +} + /// 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. diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 57373b4b130..822caccdcbe 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -1684,16 +1684,14 @@ impl PhysicalPlanner { )?, ); - // A PartialMerge runs as Single, so unpack the states it evaluates to into the - // state columns that Spark expects of it. - let mut output = Arc::clone(&aggregate); + let mut output = aggregate; let mut additional_native_plans = vec![]; if proto_mode == ProtoAggregateMode::PartialMerge { + additional_native_plans.push(Arc::clone(&output)); output = crate::execution::merge_as_partial::unpack_state_structs( output, agg.grouping_exprs.len(), )?; - additional_native_plans.push(Arc::clone(&aggregate)); } // Spark's SortAggregateExec reports its output as ordered by the grouping keys, @@ -6328,16 +6326,22 @@ mod tests { #[tokio::test] async fn partial_merge_runs_as_single_and_emits_merged_states() { - use arrow::array::Int64Array; + 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) and count(v) grouped by k, over [k, max state, count state]. + // 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 boolean = spark_expression::DataType { + type_id: 0, + type_info: None, + }; let v = Expr { expr_struct: Some(Unbound(spark_expression::UnboundReference { name: "v".to_string(), @@ -6350,7 +6354,13 @@ mod tests { plan_id: 1, children: vec![Operator { op_struct: Some(OpStruct::Scan(spark_operator::Scan { - fields: vec![int64.clone(); 3], + fields: vec![ + int64.clone(), + int64.clone(), + int64.clone(), + int64.clone(), + boolean.clone(), + ], source: String::new(), })), ..Default::default() @@ -6377,9 +6387,17 @@ mod tests { })), ..Default::default() }, + spark_expression::AggExpr { + expr_struct: Some(AggExprStruct::First(spark_expression::First { + child: Some(v.clone()), + datatype: Some(int64.clone()), + ignore_nulls: false, + })), + ..Default::default() + }, ], mode, - expr_modes: vec![partial_merge; 2], + expr_modes: vec![partial_merge; 3], initial_input_buffer_offset: 1, ..Default::default() })), @@ -6422,6 +6440,14 @@ mod tests { Some(1), ])), Arc::new(Int64Array::from(vec![1, 2, 3, 0, 4])), + Arc::new(Int64Array::from(vec![ + Some(100), + None, + Some(101), + Some(200), + None, + ])), + Arc::new(BooleanArray::from(vec![true, false, true, true, false])), ], ) .unwrap(); @@ -6440,21 +6466,18 @@ mod tests { let results = collect(native_plan.execute(0, task_ctx).unwrap()) .await .unwrap(); - let output = arrow::compute::concat_batches(&results[0].schema(), &results).unwrap(); - let column = |idx: usize| { - output - .column(idx) - .as_any() - .downcast_ref::() - .unwrap() - .clone() - }; - let (keys, maxes, counts) = (column(0), column(1), column(2)); - let mut rows: Vec<_> = (0..output.num_rows()) - .map(|row| (keys.value(row), maxes.value(row), counts.value(row))) - .collect(); - rows.sort(); - assert_eq!(vec![(1, 10, 4), (2, 5, 2), (3, 1, 4)], rows); + assert_batches_sorted_eq!( + [ + "+-------+------------+--------------+--------------------+---------------------------+", + "| col_0 | max[value] | count[count] | first[first_value] | first[first_value_is_set] |", + "+-------+------------+--------------+--------------------+---------------------------+", + "| 1 | 10 | 4 | 100 | true |", + "| 2 | 5 | 2 | 200 | true |", + "| 3 | 1 | 4 | | false |", + "+-------+------------+--------------+--------------------+---------------------------+", + ], + &results + ); } #[tokio::test] 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 f4993ff34e6..2790838a310 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -254,9 +254,7 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "1", CometConf.COMET_BATCH_SIZE.key -> "1024", - CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002", - CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true", - "spark.comet.datafusion.execution.sort_spill_reservation_bytes" -> "65536")(f) + CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002")(f) } test("PartialMerge that runs short of memory keeps one row per key (issue #6857)") { @@ -313,12 +311,13 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { val expected = (0 until 4096).map { k => Row(k, 2L, (0 until 8).map(j => padded(k + 4096 * j)).sorted) } + val localRows = rows.toDF("k", "d", "s") withTempPath { dir => val path = dir.getCanonicalPath - rows.toDF("k", "d", "s").write.parquet(path) + localRows.write.parquet(path) for (sparkPartial <- Seq(false, true)) { withDistinctMergeMemoryPressure { - val source = if (sparkPartial) rows.toDF("k", "d", "s") else spark.read.parquet(path) + val source = if (sparkPartial) localRows else spark.read.parquet(path) val df = source.groupBy("k").agg(count_distinct(col("d")), sort_array(collect_list("s"))) checkCometAnswer(df, expected)