andygrove commented on code in PR #7222:
URL: https://github.com/apache/datafusion-comet/pull/7222#discussion_r4379081041
##########
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:
Done. The guard walks the input types once against each aggregate's state
fields and returns on the first mismatch, then checks that no input columns are
left over. It builds no vectors.
##########
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:
Done. A `collect_over_batches` helper takes the planned plan and the input
batches and returns the collected output. The keyless, list-key, `PartialMerge`
and ordered `Final` tests use it now. `create_int64_datatype`,
`create_unbound_int64`, `create_max_agg` and `create_count_agg` replace the
builders the two `PartialMerge` tests each repeated.
--
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]
