comphead commented on code in PR #7222:
URL: https://github.com/apache/datafusion-comet/pull/7222#discussion_r4378972603
##########
native/core/src/execution/planner/aggregate.rs:
##########
@@ -309,6 +338,35 @@ 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<AggregateFunctionExpr>],
+ input_schema: &SchemaRef,
+) -> Result<bool, ExecutionError> {
+ 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 state_fields = aggr_expr
+ .iter()
+ .map(|expr| expr.state_fields())
+ .collect::<datafusion::common::Result<Vec<_>>>()?
+ .concat();
+ let input_types = input_schema
+ .fields()
+ .iter()
+ .skip(num_keys)
+ .map(|f| f.data_type());
+ Ok(input_types.eq(state_fields.iter().map(|f| f.data_type())))
Review Comment:
This collects every aggregate's state fields into a `Vec<Vec<_>>`,
concatenates them, and only then compares. Walking the input types once returns
on the first mismatch and builds no vectors (not compiled):
```rust
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())
```
##########
native/core/src/execution/planner.rs:
##########
@@ -4852,6 +4852,181 @@ 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::common::tree_node::{Transformed, TreeNode};
+ 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, or
on g and k, which
+ // DataFusion 55's Single would run with a GroupedHashAggregateStream
that can fail to
+ // spill (#6880).
+ 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 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<spark_expression::DataType>,
+ agg_exprs: Vec<spark_expression::AggExpr>| {
+ 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![
+ 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()
+ },
+ ],
+ )
+ };
+ for (sort_keys, input_order_mode) in [
+ (vec![0], InputOrderMode::PartiallySorted(vec![0])),
+ (vec![0, 1], InputOrderMode::Sorted),
+ ] {
+ let (_, _, planned) = planner
+ .create_plan(&max_count(&sort_keys), &mut vec![], 1)
+ .unwrap();
+ assert_eq!("ProjectionExec", planned.native_plan.name());
+ let final_agg = planned.additional_native_plans[0]
+ .downcast_ref::<AggregateExec>()
+ .unwrap();
+ assert_eq!(&AggregateMode::Final, final_agg.mode());
+ assert_eq!(&input_order_mode, 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 input: Arc<dyn ExecutionPlan> =
+ MemorySourceConfig::try_new_exec(&[vec![batch]], input_schema,
None).unwrap();
+ let native_plan = Arc::clone(&planned.native_plan)
+ .transform_up(|node| {
Review Comment:
The same `ScanExec` to `MemorySourceConfig` swap is already at lines 4529,
4632 and 4829, and the `int64`, `v`, `max` and `count` builders above repeat
the ones in `partial_merge_runs_as_single_and_emits_merged_states`. A helper
that takes the planned plan and the input batches and returns the collected
output, plus shared builders for the two PartialMerge tests, would leave this
test with only what it checks: the sort keys, the mode and input order it
expects, and the rows.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
