feat: reuse range locality for aggregation

datafusion/core/Cargo.toml

@@ -217,6 +217,10 @@ required-features = ["docs_generation"]
 harness = false
 name = "aggregate_query_sql"
 
+[[bench]]
+harness = false
+name = "range_date_trunc"
+
 [[bench]]
 harness = false
 name = "csv_load"

datafusion/core/benches/range_date_trunc.rs

@@ -0,0 +1,152 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Group hourly range partitions without moving groups across partitions.
+//! The same benchmark can run on main: only the optimized plan differs.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use criterion::{Criterion, criterion_group, criterion_main};
+use datafusion::arrow::array::{ArrayRef, Int64Array, TimestampNanosecondArray};
+use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::catalog::streaming::StreamingTable;
+use datafusion::common::{Result, ScalarValue, SplitPoint};
+use datafusion::logical_expr::col;
+use datafusion::physical_expr::{
+    Partitioning, PhysicalSortExpr, RangePartitioning, expressions::col as physical_col,
+};
+use datafusion::physical_plan::{
+    collect, execution_plan::reset_plan_states, streaming::PartitionStream,
+    test::TestPartitionStream,
+};
+use datafusion::prelude::{SessionConfig, SessionContext};
+use tokio::runtime::Runtime;
+
+const PARTITIONS: usize = 8;
+const REPETITIONS: i64 = 8;
+
+fn context(keys_per_partition: i64) -> Result<SessionContext> {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("key", DataType::Int64, false),
+        Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), false),
+        Field::new("value", DataType::Int64, false),
+    ]));
+    let start = 1_704_067_200_000_000_000_i64;
+    let hour = 3_600_000_000_000_i64;
+    let mut streams: Vec<Arc<dyn PartitionStream>> = vec![];
+    for partition in 0..PARTITIONS {
+        let mut batches = vec![];
+        let rows = keys_per_partition * REPETITIONS;
+        for offset in (0..rows).step_by(8192) {
+            let end = (offset + 8192).min(rows);
+            let columns: Vec<ArrayRef> = vec![
+                Arc::new(Int64Array::from_iter_values(
+                    (offset..end).map(|i| i / REPETITIONS),
+                )),
+                Arc::new(TimestampNanosecondArray::from_iter_values(
+                    (offset..end).map(|i| {
+                        start
+                            + partition as i64 * hour
+                            + (i % REPETITIONS) * 1_000_000_000
+                    }),
+                )),
+                Arc::new(Int64Array::from_iter_values((offset..end).map(|_| 1))),
+            ];
+            batches.push(RecordBatch::try_new(Arc::clone(&schema), columns)?);
+        }
+        streams.push(Arc::new(TestPartitionStream::new_with_batches(batches)));
+    }
+    let partitioning = Partitioning::Range(RangePartitioning::try_new(
+        [PhysicalSortExpr::new_default(physical_col("ts", &schema)?)].into(),
+        (1..PARTITIONS)
+            .map(|partition| {
+                SplitPoint::new(vec![ScalarValue::TimestampNanosecond(
+                    Some(start + partition as i64 * hour),
+                    None,
+                )])
+            })
+            .collect(),
+    )?);
+    let mut config = SessionConfig::new().with_target_partitions(PARTITIONS);
+    config
+        .options_mut()
+        .optimizer
+        .enable_round_robin_repartition = false;
+    config.options_mut().optimizer.subset_repartition_threshold = PARTITIONS;
+    let context = SessionContext::new_with_config(config);
+    context.register_table(
+        "t",
+        Arc::new(
+            StreamingTable::try_new(schema, streams)?
+                .with_sort_order(vec![
+                    col("key").sort(true, true),
+                    col("ts").sort(true, true),
+                ])
+                .with_output_partitioning(partitioning),
+        ),
+    )?;
+    Ok(context)
+}
+
+fn criterion_benchmark(criterion: &mut Criterion) {
+    let runtime = Runtime::new().unwrap();
+    for keys_per_partition in [1024, 65_536] {
+        let context = context(keys_per_partition).unwrap();
+        let plan =
+            runtime.block_on(async {
+                context
+                .sql("SELECT key, date_trunc('hour', ts) AS bucket, sum(value) AS total \
+                      FROM t GROUP BY key, bucket")
+                .await
+                .unwrap()
+                .create_physical_plan()
+                .await
+                .unwrap()
+            });
+        let results = runtime
+            .block_on(collect(Arc::clone(&plan), context.task_ctx()))
+            .unwrap();
+        assert_eq!(
+            results.iter().map(RecordBatch::num_rows).sum::<usize>(),
+            PARTITIONS * keys_per_partition as usize
+        );
+        for batch in results {
+            let sums = batch
+                .column(2)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap();
+            assert!(sums.values().iter().all(|sum| *sum == REPETITIONS));
+        }
+        criterion.bench_function(
+            &format!("range_date_trunc/{keys_per_partition}_keys_per_partition"),
+            |bencher| {
+                bencher.iter(|| {
+                    let plan = reset_plan_states(Arc::clone(&plan)).unwrap();
+                    black_box(
+                        runtime.block_on(collect(plan, context.task_ctx())).unwrap(),
+                    )
+                });
+            },
+        );
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -402,6 +402,110 @@ fn parquet_exec_with_output_partitioning(
     DataSourceExec::from_data_source(config)
 }
 
+#[test]
+fn range_date_trunc_locality_survives_aggregate_combination() -> Result<()> {
+    use arrow_schema::TimeUnit;
+    use datafusion_physical_expr::ScalarFunctionExpr;
+    use datafusion_physical_optimizer::combine_partial_final_agg::CombinePartialFinalAggregate;
+
+    let timestamp_type = DataType::Timestamp(TimeUnit::Nanosecond, None);
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "ts",
+        timestamp_type.clone(),
+        true,
+    )]));
+    let timestamp = col("ts", &schema)?;
+    let range = Partitioning::Range(RangePartitioning::try_new(
+        [PhysicalSortExpr::new_default(Arc::clone(&timestamp))].into(),
+        vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond(
+            Some(1_704_070_800_000_000_000),
+            None,
+        )])],
+    )?);
+    let scan = FileScanConfigBuilder::new(
+        ObjectStoreUrl::parse("test:///")?,
+        Arc::new(ParquetSource::new(Arc::clone(&schema))),
+    )
+    .with_file_groups(vec![
+        FileGroup::new(vec![PartitionedFile::new("p0", 100)]),
+        FileGroup::new(vec![PartitionedFile::new("p1", 100)]),
+    ])
+    .with_output_partitioning(Some(range))
+    .build();
+    let source: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(scan);
+    let bucket: Arc<dyn PhysicalExpr> = Arc::new(ScalarFunctionExpr::new(
+        "date_trunc",
+        datafusion_functions::datetime::date_trunc(),
+        vec![lit("hour"), timestamp],
+        Field::new("bucket", timestamp_type, true).into(),
+        Arc::new(ConfigOptions::default()),
+    ));
+    let group_by =
+        PhysicalGroupBy::new_single(vec![(Arc::clone(&bucket), "bucket".into())]);
+    let partial = AggregateExec::try_new(
+        AggregateMode::Partial,
+        group_by.clone(),
+        vec![],
+        vec![],
+        Arc::clone(&source),
+        Arc::clone(&schema),
+    )?;
+    let final_plan: Arc<dyn ExecutionPlan> = Arc::new(AggregateExec::try_new(
+        AggregateMode::FinalPartitioned,
+        group_by.as_final(),
+        vec![],
+        vec![],
+        Arc::new(partial),
+        schema,
+    )?);
+    let mut config = ConfigOptions::default();
+    config.execution.target_partitions = 2;
+    config.optimizer.enable_round_robin_repartition = false;
+    config.optimizer.repartition_file_scans = false;
+    let enforced = EnsureRequirements::new().optimize(final_plan, &config)?;
+    let plan_text = displayable(enforced.as_ref()).indent(true).to_string();
+    assert!(plan_text.contains("mode=FinalPartitioned"), "{plan_text}");
+    assert!(plan_text.contains("mode=Partial"), "{plan_text}");
+    assert!(!plan_text.contains("RepartitionExec"), "{plan_text}");
+    SanityCheckPlan::new().optimize(Arc::clone(&enforced), &config)?;
+
+    let combined = CombinePartialFinalAggregate::new().optimize(enforced, &config)?;
+    let combined_text = displayable(combined.as_ref()).indent(true).to_string();
+    assert!(
+        combined_text.contains("mode=SinglePartitioned"),
+        "{combined_text}"
+    );
+    let reenforced = EnsureRequirements::new().optimize(combined, &config)?;
+    assert_eq!(
+        displayable(reenforced.as_ref()).indent(true).to_string(),
+        combined_text
+    );
+    SanityCheckPlan::new().optimize(reenforced, &config)?;
+
+    // Group locality does not establish a join-compatible partition layout.
+    let join = crate::physical_optimizer::test_utils::hash_join_exec(
+        Arc::clone(&source),
+        source,
+        vec![(Arc::clone(&bucket), bucket)],
+        None,
+        &JoinType::Inner,
+    )?;
+    assert!(
+        SanityCheckPlan::new()
+            .optimize(Arc::clone(&join), &config)
+            .is_err()
+    );
+    let enforced_join = EnsureRequirements::new().optimize(join, &config)?;
+    assert!(
+        displayable(enforced_join.as_ref())
+            .indent(true)
+            .to_string()
+            .contains("partitioning=Hash")
+    );
+    SanityCheckPlan::new().optimize(enforced_join, &config)?;
+    Ok(())
+}
+
 fn range_partitioning(
     column: &str,
     split_values: impl IntoIterator<Item = i64>,

datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs

@@ -732,12 +732,11 @@ fn add_roundrobin_on_top(
     }
 }
 
-// Partial aggregates require unspecified input distribution, but their output
-// may already satisfy the final aggregate's key distribution because partial
-// aggregation preserves/projects input partitioning. Keep that reusable output
-// partitioning intact when preserve_file_partitions would otherwise insert
-// RoundRobin below the partial aggregate.
-fn partial_aggregate_output_satisfies_final_partitioning(
+// Partial aggregates require unspecified input distribution, but their input
+// may already keep every group partition-local. Preserve that layout when
+// preserve_file_partitions would otherwise insert RoundRobin below the partial
+// aggregate, and reuse it for the final aggregate.
+fn partial_aggregate_has_partition_local_groups(
     plan: &Arc<dyn ExecutionPlan>,
     allow_subset_satisfy_partitioning: bool,
 ) -> bool {
@@ -751,15 +750,23 @@ fn partial_aggregate_output_satisfies_final_partitioning(
         return false;
     }
 
-    let key_distribution = Distribution::KeyPartitioned(aggregate.output_group_expr());
-
-    plan.output_partitioning()
+    let output_distribution = Distribution::KeyPartitioned(aggregate.output_group_expr());
+    if plan
+        .output_partitioning()
         .satisfaction(
-            &key_distribution,
+            &output_distribution,
             plan.equivalence_properties(),
             allow_subset_satisfy_partitioning,
         )
         .is_satisfied()
+    {
+        return true;
+    }
+
+    aggregate.groups_are_partition_local(
+        aggregate.input().as_ref(),
+        allow_subset_satisfy_partitioning,
+    )
 }
 
 /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator
@@ -1527,7 +1534,7 @@ pub fn ensure_distribution_with_stats(
 
             let preserve_partial_aggregate_partitioning =
                 preserve_file_partition_threshold_met
-                    && partial_aggregate_output_satisfies_final_partitioning(
+                    && partial_aggregate_has_partition_local_groups(
                         &plan,
                         allow_subset_satisfy_partitioning,
                     );
@@ -1570,7 +1577,15 @@ pub fn ensure_distribution_with_stats(
                             ChildSatisfactionOptions::new()
                                 .with_allow_subset(allow_subset_satisfy_partitioning),
                         )?
-                        .is_satisfied();
+                        .is_satisfied()
+                        || plan.downcast_ref::<AggregateExec>().is_some_and(
+                            |aggregate| {
+                                aggregate.groups_are_partition_local(
+                                    child.plan.as_ref(),
+                                    allow_subset_satisfy_partitioning,
+                                )
+                            },
+                        );
                     let preserve_satisfying_file_partitioning =
                         preserve_file_partition_threshold_met
                             && !requires_grouping_id

datafusion/physical-optimizer/src/sanity_checker.rs

@@ -30,6 +30,7 @@ use datafusion_common::config::{ConfigOptions, OptimizerOptions};
 use datafusion_common::plan_err;
 use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion_physical_expr::intervals::utils::{check_support, is_datatype_supported};
+use datafusion_physical_plan::aggregates::AggregateExec;
 use datafusion_physical_plan::execution_plan::{
     Boundedness, EmissionType, InvariantLevel,
 };
@@ -167,14 +168,17 @@ pub fn check_plan_sanity(
             }
         }
 
-        if !input_distributions
-            .child_satisfaction(
-                idx,
-                child.as_ref(),
-                ChildSatisfactionOptions::new().with_allow_subset(true),
-            )?
-            .is_satisfied()
-        {
+        let distribution_satisfaction = input_distributions.child_satisfaction(
+            idx,
+            child.as_ref(),
+            ChildSatisfactionOptions::new().with_allow_subset(true),
+        )?;
+        let accepts_partition_local_groups = plan
+            .downcast_ref::<AggregateExec>()
+            .is_some_and(|aggregate| {
+                aggregate.groups_are_partition_local(child.as_ref(), true)
+            });
+        if !distribution_satisfaction.is_satisfied() && !accepts_partition_local_groups {
             let plan_str = get_plan_string(plan);
             return plan_err!(
                 "Plan: {:?} does not satisfy distribution requirements: {}. Child-{} output partitioning: {}",