feat: use range partitioning for date_trunc 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/expr-common/src/interval_arithmetic.rs

@@ -1243,6 +1243,34 @@ fn prev_value(value: ScalarValue) -> ScalarValue {
     value_transition!(MIN, false, value)
 }
 
+/// Returns the previous distinct value, or `None` for nulls, non-finite
+/// floats, type minima, and types without a discrete predecessor.
+pub fn checked_predecessor(value: &ScalarValue) -> Option<ScalarValue> {
+    checked_adjacent(value, false)
+}
+
+/// Returns the next distinct value, or `None` for nulls, non-finite floats,
+/// type maxima, and types without a discrete successor.
+pub fn checked_successor(value: &ScalarValue) -> Option<ScalarValue> {
+    checked_adjacent(value, true)
+}
+
+fn checked_adjacent(value: &ScalarValue, successor: bool) -> Option<ScalarValue> {
+    if value.is_null()
+        || matches!(value, ScalarValue::Float32(Some(v)) if !v.is_finite())
+        || matches!(value, ScalarValue::Float64(Some(v)) if !v.is_finite())
+    {
+        return None;
+    }
+
+    let adjacent = if successor {
+        next_value(value.clone())
+    } else {
+        prev_value(value.clone())
+    };
+    (!adjacent.is_null() && adjacent != *value).then_some(adjacent)
+}
+
 trait OneTrait: Sized + std::ops::Add + std::ops::Sub {
     fn one() -> Self;
 }
@@ -2261,7 +2289,8 @@ impl NullableInterval {
 mod tests {
     use crate::{
         interval_arithmetic::{
-            Interval, handle_overflow, next_value, prev_value, satisfy_greater,
+            Interval, checked_predecessor, checked_successor, handle_overflow,
+            next_value, prev_value, satisfy_greater,
         },
         operator::Operator,
     };
@@ -2358,6 +2387,36 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_checked_adjacent_values() {
+        assert_eq!(
+            checked_predecessor(&ScalarValue::Int64(Some(10))),
+            Some(ScalarValue::Int64(Some(9)))
+        );
+        assert_eq!(
+            checked_successor(&ScalarValue::Int64(Some(10))),
+            Some(ScalarValue::Int64(Some(11)))
+        );
+        assert_eq!(checked_predecessor(&ScalarValue::Int64(None)), None);
+        assert_eq!(
+            checked_predecessor(&ScalarValue::Int64(Some(i64::MIN))),
+            None
+        );
+        assert_eq!(checked_successor(&ScalarValue::Int64(Some(i64::MAX))), None);
+        assert_eq!(
+            checked_predecessor(&ScalarValue::Utf8(Some("a".into()))),
+            None
+        );
+        assert_eq!(
+            checked_successor(&ScalarValue::Float64(Some(f64::NAN))),
+            None
+        );
+        assert_eq!(
+            checked_predecessor(&ScalarValue::Float32(Some(f32::NEG_INFINITY))),
+            None
+        );
+    }
+
     #[test]
     fn test_new_interval() -> Result<()> {
         use ScalarValue::*;