From 808fd2b038e2751a9e2c4866aa00f76769a7cebf Mon Sep 17 00:00:00 2001 From: "ben.bellick" Date: Fri, 18 Sep 2026 19:22:53 +0000 Subject: [PATCH 1/2] feat: reuse range locality for aggregation --- datafusion/core/Cargo.toml | 4 + datafusion/core/benches/range_date_trunc.rs | 152 ++++++++++++++++++ .../enforce_distribution.rs | 104 ++++++++++++ .../enforce_distribution.rs | 39 +++-- .../physical-optimizer/src/sanity_checker.rs | 20 ++- .../physical-plan/src/aggregates/mod.rs | 47 ++++++ .../src/test_context/range_partitioning.rs | 6 +- .../test_files/range_sorted_time_bin_agg.slt | 81 ++++++++-- 8 files changed, 415 insertions(+), 38 deletions(-) create mode 100644 datafusion/core/benches/range_date_trunc.rs diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 222c0ec688b78..7fcf5f48e6981 100644 --- a/datafusion/core/Cargo.toml +++ b/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" diff --git a/datafusion/core/benches/range_date_trunc.rs b/datafusion/core/benches/range_date_trunc.rs new file mode 100644 index 0000000000000..4e7c3bc2f897d --- /dev/null +++ b/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 { + 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> = 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 = 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::(), + PARTITIONS * keys_per_partition as usize + ); + for batch in results { + let sums = batch + .column(2) + .as_any() + .downcast_ref::() + .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); diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 7c9d4984a6a68..8d342ca827199 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/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(×tamp))].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 = DataSourceExec::from_data_source(scan); + let bucket: Arc = 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 = 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, diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 776e3c3bf14bd..6a124d2a4c7e5 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/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, 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::().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 diff --git a/datafusion/physical-optimizer/src/sanity_checker.rs b/datafusion/physical-optimizer/src/sanity_checker.rs index 713213b70612d..0d2138014f7ee 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/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::() + .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: {}", diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 806eacb35de2d..cfd241ea400ad 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1212,6 +1212,53 @@ impl AggregateExec { &self.input } + /// Whether a candidate input keeps each group in one partition. + /// + /// Unlike co-partitioning, aggregation only needs equality groups to be + /// local. A final stage can inherit that guarantee from a partial stage + /// even when projecting the grouping expressions loses range metadata. + /// This check does not require the two stages to be combined. + pub fn groups_are_partition_local( + &self, + input: &dyn ExecutionPlan, + allow_subset: bool, + ) -> bool { + if !matches!( + self.mode, + AggregateMode::Partial + | AggregateMode::SinglePartitioned + | AggregateMode::FinalPartitioned + ) || self.group_by.is_empty() + || self.group_by.has_grouping_set() + { + return false; + } + + let required = Distribution::KeyPartitioned(self.group_by.input_exprs()); + if input + .output_partitioning() + .satisfaction(&required, input.equivalence_properties(), allow_subset) + .is_key_local() + { + return true; + } + + self.mode == AggregateMode::FinalPartitioned + && input + .downcast_ref::() + .is_some_and(|partial| { + partial.mode == AggregateMode::Partial + && datafusion_physical_expr::physical_exprs_equal( + &partial.output_group_expr(), + &self.group_by.input_exprs(), + ) + && partial.groups_are_partition_local( + partial.input.as_ref(), + allow_subset, + ) + }) + } + /// Get the input schema before any aggregates are applied pub fn input_schema(&self) -> SchemaRef { Arc::clone(&self.input_schema) diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index becde0f3286db..739534cde542d 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -344,9 +344,9 @@ type TimeBinRow = ( /// - partition 1: `[2024-01-01 01:00, 02:00)` /// /// Files are range-partitioned on `timestamp` and sorted on `(key, timestamp)`. -/// Because `date_bin(60 seconds, timestamp)` does not straddle the hour split, -/// grouping by `(key, time_bin)` is partition-disjoint. Today's planner still -/// inserts a hash shuffle; the test pins that plan so a follow-up can remove it. +/// The hour split permits partition-local `date_trunc('hour')` aggregation, +/// but `date_trunc('day')` straddles it. Aligned `date_bin` remains a negative +/// control because the function is not admitted to range analysis. pub(super) fn register_range_sorted_time_bin_table(ctx: &SessionContext) { let schema = Arc::new(Schema::new(vec![ Field::new("key", DataType::Utf8, false), diff --git a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt index 18123a492dbd6..b480db12da0ce 100644 --- a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt +++ b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt @@ -29,15 +29,10 @@ # 1. Range([timestamp]) and output_ordering=[key, timestamp] # 2. Two file_groups, so the two 60-minute streams run in parallel # -# Improvement opportunity: -# date_bin(60s) is monotonic in timestamp and the hour split is aligned to bin -# boundaries, so (key, time_bin) is partition-disjoint. Aggregation could be a -# single streaming SinglePartitioned step with no hash shuffle. -# -# Today's plan still hash-repartitions: -# Partial AggregateExec (ordering_mode=Sorted) -# -> RepartitionExec Hash([key, date_bin(...)]) -# -> FinalPartitioned AggregateExec (ordering_mode=Sorted) +# date_bin is a negative control: even aligned bins retain the hash shuffle +# because date_bin can produce NULL from non-null inputs and is not admitted +# to range analysis. Timezone-less date_trunc is admitted, but only removes +# the shuffle when every split keeps the resulting groups partition-local. statement ok set datafusion.explain.physical_plan_only = true; @@ -82,11 +77,9 @@ physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion ########## # TEST 2: Filtered time-bin aggregation. # GROUP BY keys are (key, date_bin(timestamp)). Input is sorted on those keys -# (date_bin is monotonic in timestamp) and range-partitioned so bins do not -# overlap across the two 60-minute streams. -# -# Today this is still Partial + hash RepartitionExec + Final, even though -# ordering_mode=Sorted is already recognized. +# and the split aligns with the bins in this data. The function remains +# default-denied by range analysis, so Partial + hash RepartitionExec + Final +# is retained even though ordering_mode=Sorted is recognized. ########## query TT @@ -118,7 +111,7 @@ k2 2024-01-01T01:45:00 5 ########## # TEST 3: Same aggregation without the col4 filter. The scan still has two -# 60-minute file groups, and today's plan still hash-repartitions. +# 60-minute file groups, and date_bin still requires a hash repartition. ########## query TT @@ -146,6 +139,64 @@ k2 2024-01-01T00:30:00 7 k2 2024-01-01T01:30:00 30 k2 2024-01-01T01:45:00 5 +########## +# TEST 4: date_trunc('hour') is audited and the split is aligned, so each +# group is partition-local and aggregation needs no hash shuffle. +########## + +query TT +EXPLAIN SELECT key, date_trunc('hour', timestamp) AS time_bin, sum(value) +FROM range_sorted_time_bin +WHERE col4 = 'a' +GROUP BY key, time_bin; +---- +physical_plan +01)ProjectionExec: expr=[key@0 as key, date_trunc(Utf8("hour"),range_sorted_time_bin.timestamp)@1 as time_bin, sum(range_sorted_time_bin.value)@2 as sum(range_sorted_time_bin.value)] +02)--AggregateExec: mode=SinglePartitioned, gby=[key@0 as key, date_trunc(hour, timestamp@1) as date_trunc(Utf8("hour"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +03)----FilterExec: col4@1 = a, projection=[key@0, timestamp@2, value@3] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-1.parquet]]}, projection=[key, col4, timestamp, value], output_ordering=[key@0 ASC, timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], [(1704070800000000000)], 2), file_type=parquet, predicate=col4@4 = a, pruning_predicate=col4_null_count@2 != row_count@3 AND col4_min@0 <= a AND a <= col4_max@1, required_guarantees=[col4 in (a)] + +query TPI +SELECT key, date_trunc('hour', timestamp) AS time_bin, sum(value) +FROM range_sorted_time_bin +WHERE col4 = 'a' +GROUP BY key, time_bin +ORDER BY key, time_bin; +---- +k1 2024-01-01T00:00:00 3 +k1 2024-01-01T01:00:00 30 +k2 2024-01-01T00:00:00 7 +k2 2024-01-01T01:00:00 35 + +########## +# TEST 5: date_trunc('day') straddles the hour split, so the hash shuffle +# remains. +########## + +query TT +EXPLAIN SELECT key, date_trunc('day', timestamp) AS time_bin, sum(value) +FROM range_sorted_time_bin +WHERE col4 = 'a' +GROUP BY key, time_bin; +---- +physical_plan +01)ProjectionExec: expr=[key@0 as key, date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1 as time_bin, sum(range_sorted_time_bin.value)@2 as sum(range_sorted_time_bin.value)] +02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1 as date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([key@0, date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1 ASC +04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_trunc(day, timestamp@1) as date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +05)--------FilterExec: col4@1 = a, projection=[key@0, timestamp@2, value@3] +06)----------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-1.parquet]]}, projection=[key, col4, timestamp, value], output_ordering=[key@0 ASC, timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], [(1704070800000000000)], 2), file_type=parquet, predicate=col4@4 = a, pruning_predicate=col4_null_count@2 != row_count@3 AND col4_min@0 <= a AND a <= col4_max@1, required_guarantees=[col4 in (a)] + +query TPI +SELECT key, date_trunc('day', timestamp) AS time_bin, sum(value) +FROM range_sorted_time_bin +WHERE col4 = 'a' +GROUP BY key, time_bin +ORDER BY key, time_bin; +---- +k1 2024-01-01T00:00:00 33 +k2 2024-01-01T00:00:00 42 + ########## # CLEANUP ########## From 0d691f7c1b488e8f8e51ab5a72b37980586a4c98 Mon Sep 17 00:00:00 2001 From: "ben.bellick" Date: Sun, 20 Sep 2026 18:00:35 +0000 Subject: [PATCH 2/2] refactor: consume explicit range key locality --- datafusion/physical-plan/src/aggregates/mod.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index cfd241ea400ad..a0629c017d0e1 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1234,12 +1234,11 @@ impl AggregateExec { return false; } - let required = Distribution::KeyPartitioned(self.group_by.input_exprs()); - if input - .output_partitioning() - .satisfaction(&required, input.equivalence_properties(), allow_subset) - .is_key_local() - { + if input.output_partitioning().keeps_keys_local( + &self.group_by.input_exprs(), + input.equivalence_properties(), + allow_subset, + ) { return true; }