feat: add AggregateMode::PartialReduce for tree-reduce aggregation (#20019)

DataFusion's current `AggregateMode` enum has four variants covering
three of the four cells in the input/output matrix:

| | Input: raw data          | Input: partial state |
| - | - | - |
| Output: final values | `Single` / `SinglePartitioned` | `Final` /
`FinalPartitioned` |
| Output: partial state | `Partial`                   | ??? |

This PR adds `AggregateMode::PartialReduce` to fill in the missing cell:
it takes partially-reduced values as input, and reduces them further,
but without finalizing.

This is useful because it's the key component needed to implement
distributed tree-reduction (as seen in e.g. the Scuba or Honeycomb
papers): a set of worker nodes each perform multithreaded `Partial`
aggregations, feed those into a `PartialReduce` to reduce all of this
node's values into a single row, and then a head node collects the
outputs from all nodes' `PartialReduce` to feed into a `Final`
reduction.

PR can be reviewed commit by commit: first commit is pure
refactor/simplification; most places we were matching on `AggregateMode`
we were actually just trying to either check which row of the above
table we were in, or else which column. So now we have `is_first_stage`
(tells you which column) and `is_last_stage` (tells you which row) and
we use them everywhere.

Second commit adds `PartialReduce`, and is pretty small because
`is_first_stage`/`is_last_stage` do most of the heavy lifting. It also
adds a test demonstrating a minimal Partial -> PartialReduce -> Final
tree-reduction.

Author: njsmith

datafusion/physical-optimizer/src/aggregate_statistics.rs

@@ -20,7 +20,7 @@ use datafusion_common::Result;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::scalar::ScalarValue;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
-use datafusion_physical_plan::aggregates::AggregateExec;
+use datafusion_physical_plan::aggregates::{AggregateExec, AggregateInputMode};
 use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
 use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
 use datafusion_physical_plan::udaf::{AggregateFunctionExpr, StatisticsArgs};
@@ -116,13 +116,13 @@ impl PhysicalOptimizerRule for AggregateStatistics {
 /// the `ExecutionPlan.children()` method that returns an owned reference.
 fn take_optimizable(node: &dyn ExecutionPlan) -> Option<Arc<dyn ExecutionPlan>> {
     if let Some(final_agg_exec) = node.as_any().downcast_ref::<AggregateExec>()
-        && !final_agg_exec.mode().is_first_stage()
+        && final_agg_exec.mode().input_mode() == AggregateInputMode::Partial
         && final_agg_exec.group_expr().is_empty()
     {
         let mut child = Arc::clone(final_agg_exec.input());
         loop {
             if let Some(partial_agg_exec) = child.as_any().downcast_ref::<AggregateExec>()
-                && partial_agg_exec.mode().is_first_stage()
+                && partial_agg_exec.mode().input_mode() == AggregateInputMode::Raw
                 && partial_agg_exec.group_expr().is_empty()
                 && partial_agg_exec.filter_expr().iter().all(|e| e.is_none())
             {

datafusion/physical-optimizer/src/update_aggr_exprs.rs

@@ -25,7 +25,9 @@ use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::{Result, plan_datafusion_err};
 use datafusion_physical_expr::aggregate::AggregateFunctionExpr;
 use datafusion_physical_expr::{EquivalenceProperties, PhysicalSortRequirement};
-use datafusion_physical_plan::aggregates::{AggregateExec, concat_slices};
+use datafusion_physical_plan::aggregates::{
+    AggregateExec, AggregateInputMode, concat_slices,
+};
 use datafusion_physical_plan::windows::get_ordered_partition_by_indices;
 use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 
@@ -81,7 +83,7 @@ impl PhysicalOptimizerRule for OptimizeAggregateOrder {
                 // ordering fields may be pruned out by first stage aggregates.
                 // Hence, necessary information for proper merge is added during
                 // the first stage to the state field, which the final stage uses.
-                if !aggr_exec.mode().is_first_stage() {
+                if aggr_exec.mode().input_mode() == AggregateInputMode::Partial {
                     return Ok(Transformed::no(plan));
                 }
                 let input = aggr_exec.input();