From 35792e58fc74cd07cb06d61e21c0696bfb7bc9ee Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Wed, 19 Aug 2026 15:29:09 -0400 Subject: [PATCH 1/8] test: add range-partitioned sorted time-bin aggregation coverage Pin today's Partial + hash RepartitionExec + Final plan for GROUP BY key, date_bin(timestamp) on a table that is already Range([timestamp]) and sorted on (key, timestamp), so a follow-up can remove the shuffle. Co-authored-by: Cursor --- datafusion/sqllogictest/src/test_context.rs | 8 +- .../src/test_context/range_partitioning.rs | 178 ++++++++++++++++- .../test_files/range_sorted_time_bin_agg.slt | 183 ++++++++++++++++++ 3 files changed, 365 insertions(+), 4 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 39aa2b09e5685..7252a0936afea 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -54,7 +54,9 @@ use datafusion::{ use datafusion_spark::SessionStateBuilderSpark; use crate::is_spark_path; -use range_partitioning::register_range_partitioned_table; +use range_partitioning::{ + register_metrics_range_sorted_table, register_range_partitioned_table, +}; use async_trait::async_trait; use datafusion::common::cast::as_float64_array; @@ -179,6 +181,10 @@ impl TestContext { info!("Registering range partitioned table"); register_range_partitioned_table(test_ctx.session_ctx()); } + "range_sorted_time_bin_agg.slt" => { + info!("Registering range-sorted metrics table"); + register_metrics_range_sorted_table(test_ctx.session_ctx()); + } "metadata.slt" | "arrow_field.slt" => { info!("Registering metadata table tables"); register_metadata_tables(test_ctx.session_ctx()); diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index 3cde3939f0b7c..de4f5875c8c93 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -19,9 +19,11 @@ use std::fs::{File, create_dir_all, remove_dir_all}; use std::path::Path; use std::sync::Arc; -use arrow::array::{ArrayRef, Int32Array}; +use arrow::array::{ + ArrayRef, Int32Array, Int64Array, StringArray, TimestampNanosecondArray, +}; use arrow::compute::SortOptions; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use arrow::record_batch::RecordBatch; use datafusion::catalog::streaming::StreamingTable; use datafusion::common::{ScalarValue, SplitPoint}; @@ -29,7 +31,7 @@ use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; -use datafusion::logical_expr::{Partitioning, RangePartitioning, col}; +use datafusion::logical_expr::{Partitioning, RangePartitioning, SortExpr, col}; use datafusion::parquet::arrow::ArrowWriter; use datafusion::physical_expr::{ Partitioning as PhysicalPartitioning, PhysicalSortExpr, @@ -292,3 +294,173 @@ fn range_batch(schema: SchemaRef, rows: &[(i32, i32, i32)]) -> RecordBatch { ) .expect("range batch should be valid") } + +// ============================================================================== +// Metrics table: range-partitioned on timestamp, sorted on (key, timestamp) +// ============================================================================== + +/// Unix nanoseconds for `2024-01-01 00:00:00 UTC`. +const METRICS_EPOCH_NS: i64 = 1_704_067_200_000_000_000; +const NANOS_PER_SECOND: i64 = 1_000_000_000; +const NANOS_PER_MINUTE: i64 = 60 * NANOS_PER_SECOND; + +/// Timestamp helper: minutes and seconds after `2024-01-01 00:00:00 UTC`. +fn metrics_ts(minutes: i64, seconds: i64) -> i64 { + METRICS_EPOCH_NS + minutes * NANOS_PER_MINUTE + seconds * NANOS_PER_SECOND +} + +/// Row: (key, zone, host, pod, service, timestamp_ns, value) +type MetricsRow = ( + &'static str, + &'static str, + &'static str, + &'static str, + &'static str, + i64, + i64, +); + +/// Registers `metrics_range_sorted` for time-bin aggregation plan tests. +/// +/// Two file groups, each covering a 60-minute timestamp range: +/// - partition 0: `[2024-01-01 00:00, 01:00)` +/// - 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. +pub(super) fn register_metrics_range_sorted_table(ctx: &SessionContext) { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("zone", DataType::Utf8, false), + Field::new("host", DataType::Utf8, false), + Field::new("pod", DataType::Utf8, false), + Field::new("service", DataType::Utf8, false), + Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new("value", DataType::Int64, false), + ])); + + // Each partition covers 60 minutes. The split is aligned to the 60-second + // `date_bin` used by the test query, so time bins do not straddle files. + let hour_split = metrics_ts(60, 0); + let output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("timestamp").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(hour_split), + None, + )])], + ) + .expect("metrics range partitioning should be valid"), + ); + + // Within each 60-minute file, rows are sorted by (key, timestamp). + let partitions = vec![ + vec![ + ("k1", "z1", "h1", "p1", "a", metrics_ts(0, 10), 1), + ("k1", "z1", "h1", "p1", "a", metrics_ts(0, 40), 2), + ("k1", "z1", "h1", "p1", "b", metrics_ts(1, 10), 99), + ("k2", "z1", "h1", "p1", "a", metrics_ts(30, 0), 3), + ("k2", "z1", "h1", "p1", "a", metrics_ts(30, 30), 4), + ], + vec![ + ("k1", "z1", "h1", "p1", "a", metrics_ts(60, 10), 10), + ("k1", "z1", "h1", "p1", "a", metrics_ts(60, 40), 20), + ("k2", "z1", "h1", "p1", "a", metrics_ts(90, 0), 30), + ("k2", "z1", "h1", "p1", "a", metrics_ts(105, 0), 5), + ], + ]; + + let table_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/metrics_range_sorted"); + register_metrics_listing_table( + ctx, + "metrics_range_sorted", + &table_dir, + Arc::clone(&schema), + partitions, + output_partitioning, + vec![vec![ + col("key").sort(true, true), + col("timestamp").sort(true, true), + ]], + ); +} + +fn register_metrics_listing_table( + ctx: &SessionContext, + name: &str, + table_dir: impl AsRef, + schema: SchemaRef, + partitions: Vec>, + output_partitioning: Partitioning, + file_sort_order: Vec>, +) { + let table_dir = table_dir.as_ref(); + if table_dir.exists() { + remove_dir_all(table_dir).expect("test table dir should be removable"); + } + create_dir_all(table_dir).expect("test table dir should be created"); + for (idx, rows) in partitions.into_iter().enumerate() { + let batch = metrics_batch(Arc::clone(&schema), &rows); + let file = File::create(table_dir.join(format!("part-{idx}.parquet"))) + .expect("test table parquet partition should be created"); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None) + .expect("test table parquet writer should be created"); + writer + .write(&batch) + .expect("test table parquet partition should be written"); + writer + .close() + .expect("test table parquet writer should close"); + } + + let table_path = format!( + "{}/", + table_dir + .to_str() + .expect("test table path should be valid utf8") + ); + let table_url = + ListingTableUrl::parse(&table_path).expect("test table url should parse"); + let options = ListingOptions::new(Arc::new(ParquetFormat::default())) + .with_output_partitioning(Some(output_partitioning)) + .with_file_sort_order(file_sort_order); + let config = ListingTableConfig::new(table_url) + .with_listing_options(options) + .with_schema(schema); + let table = + ListingTable::try_new(config).expect("test listing table should be valid"); + + ctx.register_table(name, Arc::new(table)) + .expect("test listing table registration should succeed"); +} + +fn metrics_batch(schema: SchemaRef, rows: &[MetricsRow]) -> RecordBatch { + RecordBatch::try_new( + schema, + vec![ + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.0))) + as ArrayRef, + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.1))) + as ArrayRef, + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.2))) + as ArrayRef, + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.3))) + as ArrayRef, + Arc::new(StringArray::from_iter_values(rows.iter().map(|row| row.4))) + as ArrayRef, + Arc::new(TimestampNanosecondArray::from_iter_values( + rows.iter().map(|row| row.5), + )) as ArrayRef, + Arc::new(Int64Array::from_iter_values(rows.iter().map(|row| row.6))) + as ArrayRef, + ], + ) + .expect("metrics batch should be valid") +} diff --git a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt new file mode 100644 index 0000000000000..43b12502f812a --- /dev/null +++ b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt @@ -0,0 +1,183 @@ +# 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 BY on a table that is: +# - range partitioned on timestamp into two 60-minute file groups +# - sorted within each file on (key, timestamp) +# +# Query: +# SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) +# FROM metrics_range_sorted +# WHERE service = 'a' +# GROUP BY key, time_bin +# +# Scan metadata already advertises: +# 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) + +statement ok +set datafusion.explain.physical_plan_only = true; + +statement ok +set datafusion.execution.collect_statistics = false; + +statement ok +set datafusion.execution.target_partitions = 2; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 2; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 1; + +statement ok +set datafusion.optimizer.enable_round_robin_repartition = false; + +statement ok +set datafusion.optimizer.enable_join_dynamic_filter_pushdown = false; + +statement ok +set datafusion.optimizer.enable_topk_dynamic_filter_pushdown = false; + +statement ok +set datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown = false; + +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +########## +# TEST 1: Scan metadata — range partitioned on timestamp, sorted on (key, timestamp), +# two parallel file groups covering 60-minute intervals. +########## + +query TT +EXPLAIN SELECT key, timestamp, value FROM metrics_range_sorted; +---- +physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet + +########## +# 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. +########## + +query TT +EXPLAIN SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) +FROM metrics_range_sorted +WHERE service = 'a' +GROUP BY key, time_bin; +---- +physical_plan +01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as time_bin, sum(metrics_range_sorted.value)@2 as sum(metrics_range_sorted.value)] +02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 ASC +04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted +05)--------FilterExec: service@1 = a, projection=[key@0, timestamp@2, value@3] +06)----------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, service, timestamp, value], output_ordering=[key@0 ASC, timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], [(1704070800000000000)], 2), file_type=parquet, predicate=service@4 = a, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= a AND a <= service_max@1, required_guarantees=[service in (a)] + +query TPI +SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) +FROM metrics_range_sorted +WHERE service = '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:30:00 7 +k2 2024-01-01T01:30:00 30 +k2 2024-01-01T01:45:00 5 + +########## +# TEST 3: Same aggregation without the service filter. The scan still has two +# 60-minute file groups, and today's plan still hash-repartitions. +########## + +query TT +EXPLAIN SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) +FROM metrics_range_sorted +GROUP BY key, time_bin; +---- +physical_plan +01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as time_bin, sum(metrics_range_sorted.value)@2 as sum(metrics_range_sorted.value)] +02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 ASC +04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted +05)--------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet + +query TPI +SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) +FROM metrics_range_sorted +GROUP BY key, time_bin +ORDER BY key, time_bin; +---- +k1 2024-01-01T00:00:00 3 +k1 2024-01-01T00:01:00 99 +k1 2024-01-01T01:00:00 30 +k2 2024-01-01T00:30:00 7 +k2 2024-01-01T01:30:00 30 +k2 2024-01-01T01:45:00 5 + +########## +# CLEANUP +########## + +# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# reset it explicitly. +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.explain.physical_plan_only; + +statement ok +reset datafusion.execution.collect_statistics; + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + +statement ok +reset datafusion.optimizer.enable_round_robin_repartition; + +statement ok +reset datafusion.optimizer.enable_join_dynamic_filter_pushdown; + +statement ok +reset datafusion.optimizer.enable_topk_dynamic_filter_pushdown; + +statement ok +reset datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown; + +statement ok +reset datafusion.execution.parquet.pushdown_filters; From 58dea6208132cdc49a33516329ae4e0a93f1b024 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Thu, 20 Aug 2026 10:27:52 -0400 Subject: [PATCH 2/8] test: generalize range-sorted time-bin aggregation test names Use generic column and table names so the coverage does not expose a metrics-specific schema. Co-authored-by: Cursor --- datafusion/sqllogictest/src/test_context.rs | 6 +- .../src/test_context/range_partitioning.rs | 62 +++++++++---------- .../test_files/range_sorted_time_bin_agg.slt | 44 ++++++------- 3 files changed, 56 insertions(+), 56 deletions(-) diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 7252a0936afea..9b97d3f59dac4 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -55,7 +55,7 @@ use datafusion_spark::SessionStateBuilderSpark; use crate::is_spark_path; use range_partitioning::{ - register_metrics_range_sorted_table, register_range_partitioned_table, + register_range_partitioned_table, register_range_sorted_time_bin_table, }; use async_trait::async_trait; @@ -182,8 +182,8 @@ impl TestContext { register_range_partitioned_table(test_ctx.session_ctx()); } "range_sorted_time_bin_agg.slt" => { - info!("Registering range-sorted metrics table"); - register_metrics_range_sorted_table(test_ctx.session_ctx()); + info!("Registering range-sorted time-bin table"); + register_range_sorted_time_bin_table(test_ctx.session_ctx()); } "metadata.slt" | "arrow_field.slt" => { info!("Registering metadata table tables"); diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index de4f5875c8c93..d1e92c68c8fb4 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -296,21 +296,21 @@ fn range_batch(schema: SchemaRef, rows: &[(i32, i32, i32)]) -> RecordBatch { } // ============================================================================== -// Metrics table: range-partitioned on timestamp, sorted on (key, timestamp) +// Time-bin table: range-partitioned on timestamp, sorted on (key, timestamp) // ============================================================================== /// Unix nanoseconds for `2024-01-01 00:00:00 UTC`. -const METRICS_EPOCH_NS: i64 = 1_704_067_200_000_000_000; +const TIME_BIN_EPOCH_NS: i64 = 1_704_067_200_000_000_000; const NANOS_PER_SECOND: i64 = 1_000_000_000; const NANOS_PER_MINUTE: i64 = 60 * NANOS_PER_SECOND; /// Timestamp helper: minutes and seconds after `2024-01-01 00:00:00 UTC`. -fn metrics_ts(minutes: i64, seconds: i64) -> i64 { - METRICS_EPOCH_NS + minutes * NANOS_PER_MINUTE + seconds * NANOS_PER_SECOND +fn time_bin_ts(minutes: i64, seconds: i64) -> i64 { + TIME_BIN_EPOCH_NS + minutes * NANOS_PER_MINUTE + seconds * NANOS_PER_SECOND } -/// Row: (key, zone, host, pod, service, timestamp_ns, value) -type MetricsRow = ( +/// Row: (key, col1, col2, col3, col4, timestamp_ns, value) +type TimeBinRow = ( &'static str, &'static str, &'static str, @@ -320,7 +320,7 @@ type MetricsRow = ( i64, ); -/// Registers `metrics_range_sorted` for time-bin aggregation plan tests. +/// Registers `range_sorted_time_bin` for time-bin aggregation plan tests. /// /// Two file groups, each covering a 60-minute timestamp range: /// - partition 0: `[2024-01-01 00:00, 01:00)` @@ -330,13 +330,13 @@ type MetricsRow = ( /// 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. -pub(super) fn register_metrics_range_sorted_table(ctx: &SessionContext) { +pub(super) fn register_range_sorted_time_bin_table(ctx: &SessionContext) { let schema = Arc::new(Schema::new(vec![ Field::new("key", DataType::Utf8, false), - Field::new("zone", DataType::Utf8, false), - Field::new("host", DataType::Utf8, false), - Field::new("pod", DataType::Utf8, false), - Field::new("service", DataType::Utf8, false), + Field::new("col1", DataType::Utf8, false), + Field::new("col2", DataType::Utf8, false), + Field::new("col3", DataType::Utf8, false), + Field::new("col4", DataType::Utf8, false), Field::new( "timestamp", DataType::Timestamp(TimeUnit::Nanosecond, None), @@ -347,7 +347,7 @@ pub(super) fn register_metrics_range_sorted_table(ctx: &SessionContext) { // Each partition covers 60 minutes. The split is aligned to the 60-second // `date_bin` used by the test query, so time bins do not straddle files. - let hour_split = metrics_ts(60, 0); + let hour_split = time_bin_ts(60, 0); let output_partitioning = Partitioning::Range( RangePartitioning::try_new( vec![col("timestamp").sort(true, true)], @@ -356,31 +356,31 @@ pub(super) fn register_metrics_range_sorted_table(ctx: &SessionContext) { None, )])], ) - .expect("metrics range partitioning should be valid"), + .expect("time-bin range partitioning should be valid"), ); // Within each 60-minute file, rows are sorted by (key, timestamp). let partitions = vec![ vec![ - ("k1", "z1", "h1", "p1", "a", metrics_ts(0, 10), 1), - ("k1", "z1", "h1", "p1", "a", metrics_ts(0, 40), 2), - ("k1", "z1", "h1", "p1", "b", metrics_ts(1, 10), 99), - ("k2", "z1", "h1", "p1", "a", metrics_ts(30, 0), 3), - ("k2", "z1", "h1", "p1", "a", metrics_ts(30, 30), 4), + ("k1", "x", "y", "z", "a", time_bin_ts(0, 10), 1), + ("k1", "x", "y", "z", "a", time_bin_ts(0, 40), 2), + ("k1", "x", "y", "z", "b", time_bin_ts(1, 10), 99), + ("k2", "x", "y", "z", "a", time_bin_ts(30, 0), 3), + ("k2", "x", "y", "z", "a", time_bin_ts(30, 30), 4), ], vec![ - ("k1", "z1", "h1", "p1", "a", metrics_ts(60, 10), 10), - ("k1", "z1", "h1", "p1", "a", metrics_ts(60, 40), 20), - ("k2", "z1", "h1", "p1", "a", metrics_ts(90, 0), 30), - ("k2", "z1", "h1", "p1", "a", metrics_ts(105, 0), 5), + ("k1", "x", "y", "z", "a", time_bin_ts(60, 10), 10), + ("k1", "x", "y", "z", "a", time_bin_ts(60, 40), 20), + ("k2", "x", "y", "z", "a", time_bin_ts(90, 0), 30), + ("k2", "x", "y", "z", "a", time_bin_ts(105, 0), 5), ], ]; let table_dir = Path::new(env!("CARGO_MANIFEST_DIR")) - .join("test_files/scratch_range_partitioning/metrics_range_sorted"); - register_metrics_listing_table( + .join("test_files/scratch_range_partitioning/range_sorted_time_bin"); + register_time_bin_listing_table( ctx, - "metrics_range_sorted", + "range_sorted_time_bin", &table_dir, Arc::clone(&schema), partitions, @@ -392,12 +392,12 @@ pub(super) fn register_metrics_range_sorted_table(ctx: &SessionContext) { ); } -fn register_metrics_listing_table( +fn register_time_bin_listing_table( ctx: &SessionContext, name: &str, table_dir: impl AsRef, schema: SchemaRef, - partitions: Vec>, + partitions: Vec>, output_partitioning: Partitioning, file_sort_order: Vec>, ) { @@ -407,7 +407,7 @@ fn register_metrics_listing_table( } create_dir_all(table_dir).expect("test table dir should be created"); for (idx, rows) in partitions.into_iter().enumerate() { - let batch = metrics_batch(Arc::clone(&schema), &rows); + let batch = time_bin_batch(Arc::clone(&schema), &rows); let file = File::create(table_dir.join(format!("part-{idx}.parquet"))) .expect("test table parquet partition should be created"); let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None) @@ -441,7 +441,7 @@ fn register_metrics_listing_table( .expect("test listing table registration should succeed"); } -fn metrics_batch(schema: SchemaRef, rows: &[MetricsRow]) -> RecordBatch { +fn time_bin_batch(schema: SchemaRef, rows: &[TimeBinRow]) -> RecordBatch { RecordBatch::try_new( schema, vec![ @@ -462,5 +462,5 @@ fn metrics_batch(schema: SchemaRef, rows: &[MetricsRow]) -> RecordBatch { as ArrayRef, ], ) - .expect("metrics batch should be valid") + .expect("time-bin batch should be valid") } 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 43b12502f812a..18123a492dbd6 100644 --- a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt +++ b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt @@ -21,8 +21,8 @@ # # Query: # SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) -# FROM metrics_range_sorted -# WHERE service = 'a' +# FROM range_sorted_time_bin +# WHERE col4 = 'a' # GROUP BY key, time_bin # # Scan metadata already advertises: @@ -75,9 +75,9 @@ set datafusion.execution.parquet.pushdown_filters = false; ########## query TT -EXPLAIN SELECT key, timestamp, value FROM metrics_range_sorted; +EXPLAIN SELECT key, timestamp, value FROM range_sorted_time_bin; ---- -physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet +physical_plan 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, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet ########## # TEST 2: Filtered time-bin aggregation. @@ -91,22 +91,22 @@ physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion query TT EXPLAIN SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) -FROM metrics_range_sorted -WHERE service = 'a' +FROM range_sorted_time_bin +WHERE col4 = 'a' GROUP BY key, time_bin; ---- physical_plan -01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as time_bin, sum(metrics_range_sorted.value)@2 as sum(metrics_range_sorted.value)] -02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted -03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 ASC -04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted -05)--------FilterExec: service@1 = a, projection=[key@0, timestamp@2, value@3] -06)----------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, service, timestamp, value], output_ordering=[key@0 ASC, timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], [(1704070800000000000)], 2), file_type=parquet, predicate=service@4 = a, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= a AND a <= service_max@1, required_guarantees=[service in (a)] +01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),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_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 ASC +04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),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_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) -FROM metrics_range_sorted -WHERE service = 'a' +FROM range_sorted_time_bin +WHERE col4 = 'a' GROUP BY key, time_bin ORDER BY key, time_bin; ---- @@ -117,25 +117,25 @@ k2 2024-01-01T01:30:00 30 k2 2024-01-01T01:45:00 5 ########## -# TEST 3: Same aggregation without the service filter. The scan still has two +# TEST 3: Same aggregation without the col4 filter. The scan still has two # 60-minute file groups, and today's plan still hash-repartitions. ########## query TT EXPLAIN SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) -FROM metrics_range_sorted +FROM range_sorted_time_bin GROUP BY key, time_bin; ---- physical_plan -01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as time_bin, sum(metrics_range_sorted.value)@2 as sum(metrics_range_sorted.value)] -02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted -03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)@1 ASC -04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),metrics_range_sorted.timestamp)], aggr=[sum(metrics_range_sorted.value)], ordering_mode=Sorted -05)--------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/metrics_range_sorted/part-1.parquet]]}, projection=[key, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet +01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),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_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 ASC +04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +05)--------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, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet query TPI SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) -FROM metrics_range_sorted +FROM range_sorted_time_bin GROUP BY key, time_bin ORDER BY key, time_bin; ---- From 600609629edb3033ca9f688070599ec9f7940d11 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Thu, 20 Aug 2026 11:34:12 -0400 Subject: [PATCH 3/8] test: share parquet listing-table helper for range tests Accept RecordBatches and optional file sort order so the time-bin table reuses the same registration path. Co-authored-by: Cursor --- .../src/test_context/range_partitioning.rs | 96 +++++++------------ 1 file changed, 34 insertions(+), 62 deletions(-) diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index d1e92c68c8fb4..becde0f3286db 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -97,8 +97,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { "range_partitioned", &range_table_dir, Arc::clone(&schema), - RANGE_PARTITIONS, + range_batches(&schema, RANGE_PARTITIONS), output_partitioning, + None, ); register_unbounded_range_stream_table( @@ -134,8 +135,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { Path::new(env!("CARGO_MANIFEST_DIR")) .join("test_files/scratch_range_partitioning/range_partitioned_shifted"), Arc::clone(&schema), - SHIFTED_RANGE_PARTITIONS, + range_batches(&schema, SHIFTED_RANGE_PARTITIONS), shifted_output_partitioning, + None, ); // Same rows as `range_partitioned` but split into only three range @@ -158,8 +160,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { Path::new(env!("CARGO_MANIFEST_DIR")) .join("test_files/scratch_range_partitioning/range_partitioned_narrow"), Arc::clone(&schema), - NARROW_RANGE_PARTITIONS, + range_batches(&schema, NARROW_RANGE_PARTITIONS), narrow_output_partitioning, + None, ); let sparse_output_partitioning = Partitioning::Range( @@ -180,8 +183,9 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { Path::new(env!("CARGO_MANIFEST_DIR")) .join("test_files/scratch_range_partitioning/range_partitioned_sparse"), Arc::clone(&schema), - SPARSE_RANGE_PARTITIONS, + range_batches(&schema, SPARSE_RANGE_PARTITIONS), sparse_output_partitioning, + None, ); } @@ -190,16 +194,16 @@ fn register_parquet_listing_table( name: &str, table_dir: impl AsRef, schema: SchemaRef, - partitions: impl IntoIterator, + batches: Vec, output_partitioning: Partitioning, + file_sort_order: Option>>, ) { let table_dir = table_dir.as_ref(); if table_dir.exists() { remove_dir_all(table_dir).expect("test table dir should be removable"); } create_dir_all(table_dir).expect("test table dir should be created"); - for (idx, rows) in partitions.into_iter().enumerate() { - let batch = range_batch(Arc::clone(&schema), rows); + for (idx, batch) in batches.into_iter().enumerate() { let file = File::create(table_dir.join(format!("part-{idx}.parquet"))) .expect("test table parquet partition should be created"); let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None) @@ -220,8 +224,11 @@ fn register_parquet_listing_table( ); let table_url = ListingTableUrl::parse(&table_path).expect("test table url should parse"); - let options = ListingOptions::new(Arc::new(ParquetFormat::default())) + let mut options = ListingOptions::new(Arc::new(ParquetFormat::default())) .with_output_partitioning(Some(output_partitioning)); + if let Some(file_sort_order) = file_sort_order { + options = options.with_file_sort_order(file_sort_order); + } let config = ListingTableConfig::new(table_url) .with_listing_options(options) .with_schema(schema); @@ -280,6 +287,16 @@ fn range_stream_partition( )])) } +fn range_batches( + schema: &SchemaRef, + partitions: impl IntoIterator, +) -> Vec { + partitions + .into_iter() + .map(|rows| range_batch(Arc::clone(schema), rows)) + .collect() +} + fn range_batch(schema: SchemaRef, rows: &[(i32, i32, i32)]) -> RecordBatch { RecordBatch::try_new( schema, @@ -360,7 +377,7 @@ pub(super) fn register_range_sorted_time_bin_table(ctx: &SessionContext) { ); // Within each 60-minute file, rows are sorted by (key, timestamp). - let partitions = vec![ + let partitions = [ vec![ ("k1", "x", "y", "z", "a", time_bin_ts(0, 10), 1), ("k1", "x", "y", "z", "a", time_bin_ts(0, 40), 2), @@ -378,69 +395,24 @@ pub(super) fn register_range_sorted_time_bin_table(ctx: &SessionContext) { let table_dir = Path::new(env!("CARGO_MANIFEST_DIR")) .join("test_files/scratch_range_partitioning/range_sorted_time_bin"); - register_time_bin_listing_table( + let batches = partitions + .iter() + .map(|rows| time_bin_batch(Arc::clone(&schema), rows)) + .collect(); + register_parquet_listing_table( ctx, "range_sorted_time_bin", &table_dir, Arc::clone(&schema), - partitions, + batches, output_partitioning, - vec![vec![ + Some(vec![vec![ col("key").sort(true, true), col("timestamp").sort(true, true), - ]], + ]]), ); } -fn register_time_bin_listing_table( - ctx: &SessionContext, - name: &str, - table_dir: impl AsRef, - schema: SchemaRef, - partitions: Vec>, - output_partitioning: Partitioning, - file_sort_order: Vec>, -) { - let table_dir = table_dir.as_ref(); - if table_dir.exists() { - remove_dir_all(table_dir).expect("test table dir should be removable"); - } - create_dir_all(table_dir).expect("test table dir should be created"); - for (idx, rows) in partitions.into_iter().enumerate() { - let batch = time_bin_batch(Arc::clone(&schema), &rows); - let file = File::create(table_dir.join(format!("part-{idx}.parquet"))) - .expect("test table parquet partition should be created"); - let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None) - .expect("test table parquet writer should be created"); - writer - .write(&batch) - .expect("test table parquet partition should be written"); - writer - .close() - .expect("test table parquet writer should close"); - } - - let table_path = format!( - "{}/", - table_dir - .to_str() - .expect("test table path should be valid utf8") - ); - let table_url = - ListingTableUrl::parse(&table_path).expect("test table url should parse"); - let options = ListingOptions::new(Arc::new(ParquetFormat::default())) - .with_output_partitioning(Some(output_partitioning)) - .with_file_sort_order(file_sort_order); - let config = ListingTableConfig::new(table_url) - .with_listing_options(options) - .with_schema(schema); - let table = - ListingTable::try_new(config).expect("test listing table should be valid"); - - ctx.register_table(name, Arc::new(table)) - .expect("test listing table registration should succeed"); -} - fn time_bin_batch(schema: SchemaRef, rows: &[TimeBinRow]) -> RecordBatch { RecordBatch::try_new( schema, From b26846d8796ed6e5231f2ba53a50ced4e4ebacdd Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Wed, 19 Aug 2026 15:54:45 -0400 Subject: [PATCH 4/8] feat: skip hash shuffle for date_bin/date_trunc on Range([timestamp]) Treat Range([x]) as a subset of KeyPartitioned([..., f(x), ...]) when f is monotonic and bins do not straddle split points, so GROUP BY key, date_bin/date_trunc can stream without a hash RepartitionExec. Co-authored-by: Cursor --- datafusion/physical-expr/src/partitioning.rs | 438 +++++++++++++++++- .../src/test_context/range_partitioning.rs | 8 +- .../test_files/range_sorted_time_bin_agg.slt | 97 +++- 3 files changed, 495 insertions(+), 48 deletions(-) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 98f082f7256db..b493a9b10d84e 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -17,14 +17,23 @@ //! [`Partitioning`] and [`Distribution`] for `ExecutionPlans` +use crate::expressions::{Column, Literal, UnKnownColumn}; +use crate::utils::collect_columns; use crate::{ EquivalenceProperties, PhysicalExpr, equivalence::ProjectionMapping, - expressions::UnKnownColumn, physical_exprs_contains, physical_exprs_equal, + physical_exprs_contains, physical_exprs_equal, }; pub use datafusion_common::SplitPoint; -use datafusion_common::{Result, validate_range_split_points}; +use datafusion_common::{Result, ScalarValue, validate_range_split_points}; +use datafusion_expr::ColumnarValue; +use datafusion_expr::interval_arithmetic::Interval; +use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_physical_expr_common::physical_expr::format_physical_expr_list; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; + +use arrow::array::new_null_array; +use arrow::datatypes::Schema; +use arrow::record_batch::RecordBatch; #[cfg(feature = "proto")] use datafusion_physical_expr_common::sort_expr::{ sort_exprs_try_from_proto, sort_exprs_try_to_proto, @@ -252,25 +261,51 @@ impl RangePartitioning { /// /// Returns `None` if any range key cannot be projected or if projection /// collapses distinct range keys into duplicate output expressions. + /// + /// A range key that is not emitted as-is can still be projected when the + /// mapping contains a monotonic function of that key (for example + /// `date_bin(interval, timestamp)` or `date_trunc(unit, timestamp)` while + /// range-partitioned on `timestamp`). + /// Adjacent partitions stay disjoint only when evaluating the function at + /// each split point and its predecessor yields different values, so bins + /// do not straddle file groups. fn project( &self, mapping: &ProjectionMapping, input_eq_properties: &EquivalenceProperties, ) -> Option { - let exprs = self - .ordering - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .collect::>(); - let projected_exprs = input_eq_properties - .project_expressions(&exprs, mapping) - .collect::>>()?; - let sort_exprs = self - .ordering - .iter() - .zip(projected_exprs) - .map(|(sort_expr, expr)| PhysicalSortExpr::new(expr, sort_expr.options)) - .collect::>(); + let mut split_points = self.split_points.clone(); + let mut sort_exprs = Vec::with_capacity(self.ordering.len()); + for (key_idx, sort_expr) in self.ordering.iter().enumerate() { + if let Some(projected) = + input_eq_properties.project_expr(&sort_expr.expr, mapping) + { + sort_exprs.push(PhysicalSortExpr::new(projected, sort_expr.options)); + continue; + } + + let (target, source) = + monotonic_range_key_projection(sort_expr, mapping, input_eq_properties)?; + if !monotonic_fn_keeps_partitions_disjoint( + &source, + &sort_expr.expr, + &split_points, + key_idx, + input_eq_properties.schema(), + ) { + return None; + } + if let Some(updated) = project_split_points_through_fn( + &source, + &sort_expr.expr, + &split_points, + key_idx, + input_eq_properties.schema(), + ) { + split_points = updated; + } + sort_exprs.push(PhysicalSortExpr::new(target, sort_expr.options)); + } let ordering = LexOrdering::new(sort_exprs)?; if ordering.len() != self.ordering.len() { return None; @@ -278,9 +313,210 @@ impl RangePartitioning { Some(Self { ordering, - split_points: self.split_points.clone(), + split_points, + }) + } +} + +/// Finds a projection mapping whose source is a monotonic function of `sort_expr`. +fn monotonic_range_key_projection( + sort_expr: &PhysicalSortExpr, + mapping: &ProjectionMapping, + eq_properties: &EquivalenceProperties, +) -> Option<(Arc, Arc)> { + mapping.iter().find_map(|(source, targets)| { + is_order_preserving_function_of(source, &sort_expr.expr, eq_properties.schema()) + .then(|| (Arc::clone(&targets.first().0), Arc::clone(source))) + }) +} + +/// Returns true when `expr` is a (possibly non-strict) monotonic function of +/// `range_key` plus literals, such as `date_bin(interval, timestamp)` or +/// `date_trunc(unit, timestamp)`. +fn is_order_preserving_function_of( + expr: &Arc, + range_key: &Arc, + schema: &arrow::datatypes::SchemaRef, +) -> bool { + if expr.eq(range_key) { + return false; + } + let expr_cols = collect_columns(expr); + let key_cols = collect_columns(range_key); + if key_cols.is_empty() || expr_cols != key_cols { + return false; + } + let Ok(child_properties) = expr + .children() + .iter() + .map(|child| function_child_properties(child, range_key, schema)) + .collect::>>() + else { + return false; + }; + matches!( + expr.get_properties(&child_properties) + .map(|properties| properties.sort_properties), + Ok(SortProperties::Ordered(_)) + ) +} + +fn function_child_properties( + child: &Arc, + range_key: &Arc, + schema: &arrow::datatypes::SchemaRef, +) -> Result { + if child.eq(range_key) { + let data_type = child.data_type(schema)?; + return Ok(ExprProperties { + sort_properties: SortProperties::Ordered(Default::default()), + range: Interval::make_unbounded(&data_type)?, + preserves_lex_ordering: true, + strictly_order_preserving: true, + }); + } + if child.downcast_ref::().is_some() { + return Ok(ExprProperties { + sort_properties: SortProperties::Singleton, + range: Interval::make_unbounded(&child.data_type(schema)?)?, + preserves_lex_ordering: true, + strictly_order_preserving: true, + }); + } + Ok(ExprProperties::new_unknown()) +} + +/// Adjacent range partitions remain disjoint on `fn_expr` when the function +/// value at each split differs from the value immediately below the split. +fn monotonic_fn_keeps_partitions_disjoint( + fn_expr: &Arc, + range_key: &Arc, + split_points: &[SplitPoint], + key_idx: usize, + schema: &arrow::datatypes::SchemaRef, +) -> bool { + split_points.iter().all(|split_point| { + let Some(split_value) = split_point.values().get(key_idx) else { + return false; + }; + let Some(predecessor) = scalar_predecessor(split_value) else { + return false; + }; + let Some(at_split) = + evaluate_expr_on_key(fn_expr, range_key, split_value, schema) + else { + return false; + }; + let Some(below_split) = + evaluate_expr_on_key(fn_expr, range_key, &predecessor, schema) + else { + return false; + }; + at_split != below_split + }) +} + +fn project_split_points_through_fn( + fn_expr: &Arc, + range_key: &Arc, + split_points: &[SplitPoint], + key_idx: usize, + schema: &arrow::datatypes::SchemaRef, +) -> Option> { + split_points + .iter() + .map(|split_point| { + let split_value = split_point.values().get(key_idx)?; + let projected = + evaluate_expr_on_key(fn_expr, range_key, split_value, schema)?; + let mut values = split_point.values().to_vec(); + values[key_idx] = projected; + Some(SplitPoint::new(values)) }) + .collect() +} + +fn evaluate_expr_on_key( + expr: &Arc, + range_key: &Arc, + value: &ScalarValue, + schema: &Schema, +) -> Option { + let column = range_key.downcast_ref::()?; + // The table schema may mark columns non-nullable. Build a 1-row batch with + // nullable fields so unused columns can be null while still evaluating `expr`. + let nullable_schema = Arc::new(Schema::new( + schema + .fields() + .iter() + .map(|field| field.as_ref().clone().with_nullable(true)) + .collect::>(), + )); + let arrays = nullable_schema + .fields() + .iter() + .enumerate() + .map(|(idx, field)| { + if idx == column.index() { + value.to_array_of_size(1).ok() + } else { + Some(new_null_array(field.data_type(), 1)) + } + }) + .collect::>>()?; + let batch = RecordBatch::try_new(nullable_schema, arrays).ok()?; + match expr.evaluate(&batch).ok()? { + ColumnarValue::Scalar(scalar) => Some(scalar), + ColumnarValue::Array(array) => ScalarValue::try_from_array(&array, 0).ok(), + } +} + +fn scalar_predecessor(value: &ScalarValue) -> Option { + match value { + ScalarValue::TimestampNanosecond(Some(v), tz) => Some( + ScalarValue::TimestampNanosecond(Some(v.checked_sub(1)?), tz.clone()), + ), + ScalarValue::TimestampMicrosecond(Some(v), tz) => Some( + ScalarValue::TimestampMicrosecond(Some(v.checked_sub(1)?), tz.clone()), + ), + ScalarValue::TimestampMillisecond(Some(v), tz) => Some( + ScalarValue::TimestampMillisecond(Some(v.checked_sub(1)?), tz.clone()), + ), + ScalarValue::TimestampSecond(Some(v), tz) => Some(ScalarValue::TimestampSecond( + Some(v.checked_sub(1)?), + tz.clone(), + )), + ScalarValue::Int64(Some(v)) => Some(ScalarValue::Int64(Some(v.checked_sub(1)?))), + ScalarValue::Int32(Some(v)) => Some(ScalarValue::Int32(Some(v.checked_sub(1)?))), + _ => None, + } +} + +/// Range([x]) satisfies grouping by `(..., f(x), ...)` when `f` is monotonic in +/// `x` and adjacent partitions do not share `f` values (bins do not straddle +/// split points). That makes `(key, date_bin(timestamp))` and +/// `(key, date_trunc(timestamp))` partition-disjoint when the table is +/// range-partitioned on `timestamp` and the split is aligned to the bin. +fn range_monotonic_fn_satisfies_keys( + range: &RangePartitioning, + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, +) -> bool { + if range.ordering().len() != 1 { + return false; } + let range_key = &range.ordering()[0].expr; + let schema = eq_properties.schema(); + required_exprs.iter().any(|required| { + is_order_preserving_function_of(required, range_key, schema) + && monotonic_fn_keeps_partitions_disjoint( + required, + range_key, + range.split_points(), + 0, + schema, + ) + }) } impl Display for RangePartitioning { @@ -433,12 +669,24 @@ impl Partitioning { .iter() .map(|sort_expr| Arc::clone(&sort_expr.expr)) .collect::>(); - Self::key_satisfaction( + let satisfaction = Self::key_satisfaction( &partition_exprs, required_exprs, eq_properties, allow_subset, - ) + ); + if satisfaction == PartitioningSatisfaction::NotSatisfied + && allow_subset + && range_monotonic_fn_satisfies_keys( + range, + required_exprs, + eq_properties, + ) + { + PartitioningSatisfaction::Subset + } else { + satisfaction + } } Partitioning::RoundRobinBatch(_) | Partitioning::UnknownPartitioning(_) => { @@ -746,11 +994,13 @@ impl Display for Distribution { mod tests { use super::*; - use crate::expressions::Column; + use crate::ScalarFunctionExpr; + use crate::expressions::{Column, Literal}; use crate::projection::ProjectionTargets; use arrow::compute::SortOptions; - use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; + use datafusion_common::config::ConfigOptions; use datafusion_common::{Result, ScalarValue}; struct PartitioningTestFixture { @@ -1227,6 +1477,152 @@ mod tests { Ok(()) } + fn date_bin_of( + timestamp: Arc, + stride_ns: i64, + ) -> Arc { + datetime_fn( + "date_bin", + datafusion_functions::datetime::date_bin(), + vec![ + Arc::new(Literal::new(ScalarValue::new_interval_mdn(0, 0, stride_ns))), + timestamp, + ], + ) + } + + fn date_trunc_of( + timestamp: Arc, + precision: &str, + ) -> Arc { + datetime_fn( + "date_trunc", + datafusion_functions::datetime::date_trunc(), + vec![ + Arc::new(Literal::new(ScalarValue::Utf8(Some(precision.to_string())))), + timestamp, + ], + ) + } + + fn datetime_fn( + name: &str, + fun: Arc, + args: Vec>, + ) -> Arc { + Arc::new(ScalarFunctionExpr::new( + name, + fun, + args, + Field::new( + "time_bin", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + ) + .into(), + Arc::new(ConfigOptions::default()), + )) + } + + fn ts_ns_split(ns: i64) -> SplitPoint { + SplitPoint::new(vec![ScalarValue::TimestampNanosecond(Some(ns), None)]) + } + + #[test] + fn range_partitioning_satisfies_monotonic_date_bin_grouping() -> Result<()> { + let fixture = PartitioningTestFixture::new(vec![ + ("key", DataType::Utf8), + ("timestamp", DataType::Timestamp(TimeUnit::Nanosecond, None)), + ])?; + // 2024-01-01T01:00:00, aligned to a 60-second date_bin. + let hour_ns = 1_704_070_800_000_000_000i64; + let aligned = fixture.range_partitioning([1], vec![ts_ns_split(hour_ns)]); + let unaligned = + fixture.range_partitioning([1], vec![ts_ns_split(hour_ns + 30_000_000_000)]); + + let required = Distribution::KeyPartitioned(vec![ + fixture.col(0), + date_bin_of(fixture.col(1), 60_000_000_000), + ]); + + assert_satisfaction( + "aligned hour split: Range(timestamp) subset-satisfies GROUP BY (key, date_bin(60s, timestamp))", + &aligned, + &required, + &fixture.eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + assert_satisfaction( + "unaligned split does not satisfy date_bin grouping", + &unaligned, + &required, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + let trunc_hour = Distribution::KeyPartitioned(vec![ + fixture.col(0), + date_trunc_of(fixture.col(1), "hour"), + ]); + assert_satisfaction( + "aligned hour split: Range(timestamp) subset-satisfies GROUP BY (key, date_trunc(hour, timestamp))", + &aligned, + &trunc_hour, + &fixture.eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + + let trunc_day = Distribution::KeyPartitioned(vec![ + fixture.col(0), + date_trunc_of(fixture.col(1), "day"), + ]); + assert_satisfaction( + "hour split straddles date_trunc(day) bins", + &aligned, + &trunc_day, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + Ok(()) + } + + #[test] + fn test_range_partitioning_project_through_date_bin() -> Result<()> { + let fixture = PartitioningTestFixture::new(vec![( + "timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, None), + )])?; + let hour_ns = 1_704_070_800_000_000_000i64; + let date_bin = date_bin_of(fixture.col(0), 60_000_000_000); + let target: Arc = Arc::new(Column::new("time_bin", 0)); + let mapping = ProjectionMapping::from_iter([( + Arc::clone(&date_bin), + ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), + )]); + + let aligned = fixture.range_partitioning([0], vec![ts_ns_split(hour_ns)]); + let projected = aligned.project(&mapping, &fixture.eq_properties); + assert_eq!( + projected.to_string(), + "Range([time_bin@0 ASC], [(1704070800000000000)], 2)" + ); + + let unaligned = + fixture.range_partitioning([0], vec![ts_ns_split(hour_ns + 30_000_000_000)]); + let projected = unaligned.project(&mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); + + Ok(()) + } + #[test] fn range_partitioning_key_distribution_satisfaction() -> Result<()> { let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index becde0f3286db..46308fae284b8 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -344,9 +344,11 @@ 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. +/// Because `date_bin(60 seconds, timestamp)` and `date_trunc('hour', timestamp)` +/// do not straddle the hour split, grouping by `(key, time_bin)` is +/// partition-disjoint and aggregation can run in one streaming step. Bins that +/// do straddle the split (for example `date_trunc('day', timestamp)`) still +/// require a hash shuffle. 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..cfe4f381c1f0d 100644 --- a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt +++ b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt @@ -25,19 +25,16 @@ # WHERE col4 = 'a' # GROUP BY key, time_bin # -# Scan metadata already advertises: +# Scan metadata advertises: # 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. +# date_bin(60s) and date_trunc('hour') are monotonic in timestamp and the hour +# split is aligned to those bins, so (key, time_bin) is partition-disjoint. +# Aggregation is one 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_trunc('day') bins straddle the hour split, so that query still +# hash-repartitions. statement ok set datafusion.explain.physical_plan_only = true; @@ -83,10 +80,8 @@ 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. +# overlap across the two 60-minute streams. Aggregation is one streaming +# SinglePartitioned step with no hash shuffle. ########## query TT @@ -97,11 +92,9 @@ GROUP BY key, time_bin; ---- physical_plan 01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),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_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted -03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 ASC -04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),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)] +02)--AggregateExec: mode=SinglePartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),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_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) @@ -117,8 +110,8 @@ k2 2024-01-01T01:30:00 30 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. +# TEST 3: Same aggregation without the col4 filter, still one streaming step +# across the two 60-minute file groups. ########## query TT @@ -128,10 +121,8 @@ GROUP BY key, time_bin; ---- physical_plan 01)ProjectionExec: expr=[key@0 as key, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),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_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted -03)----RepartitionExec: partitioning=Hash([key@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)@1 ASC -04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted -05)--------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, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet +02)--AggregateExec: mode=SinglePartitioned, gby=[key@0 as key, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }, timestamp@1) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 60000000000 }"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +03)----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, timestamp, value], output_ordering=[key@0 ASC, timestamp@1 ASC], output_partitioning=Range([timestamp@1 ASC], [(1704070800000000000)], 2), file_type=parquet query TPI SELECT key, date_bin(INTERVAL '60 seconds', timestamp) AS time_bin, sum(value) @@ -146,6 +137,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 aligned to the hour split, so the same +# SinglePartitioned streaming plan applies. +########## + +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') bins straddle the hour split (both file groups are +# 2024-01-01), so grouping is not partition-disjoint and a 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 f3a64ff25433dc6ae6628df1fc92cbab30cb86e6 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Thu, 20 Aug 2026 12:42:39 -0400 Subject: [PATCH 5/8] refactor: reuse existing monotonicity helpers for range projection Drop the local order-preserving walk and scalar predecessor so date_bin/date_trunc range projection uses EquivalenceProperties and interval_arithmetic instead of duplicating them. Co-authored-by: Cursor --- .../expr-common/src/interval_arithmetic.rs | 14 ++ .../src/equivalence/properties/mod.rs | 25 +++ datafusion/physical-expr/src/partitioning.rs | 153 ++++-------------- 3 files changed, 66 insertions(+), 126 deletions(-) diff --git a/datafusion/expr-common/src/interval_arithmetic.rs b/datafusion/expr-common/src/interval_arithmetic.rs index 68541e1e6b32c..3f476bb3791a5 100644 --- a/datafusion/expr-common/src/interval_arithmetic.rs +++ b/datafusion/expr-common/src/interval_arithmetic.rs @@ -1243,6 +1243,20 @@ fn prev_value(value: ScalarValue) -> ScalarValue { value_transition!(MIN, false, value) } +/// Returns the previous distinct value of `value`, or `None` if `value` is +/// null, already at the type minimum, or a type that has no predecessor. +pub fn checked_predecessor(value: &ScalarValue) -> Option { + if value.is_null() { + return None; + } + let predecessor = prev_value(value.clone()); + if predecessor.is_null() || predecessor == *value { + None + } else { + Some(predecessor) + } +} + trait OneTrait: Sized + std::ops::Add + std::ops::Sub { fn one() -> Self; } diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 54269e07f9309..29a9c7ea0e83e 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -1300,6 +1300,31 @@ impl EquivalenceProperties { .unwrap_or_else(|_| ExprProperties::new_unknown()) } + /// Returns true when `expr` is a (possibly non-strict) monotonic function of + /// `range_key` plus literals, such as `date_bin(interval, timestamp)` or + /// `date_trunc(unit, timestamp)`. + /// + /// The identity `expr == range_key` returns false so callers can treat "emit + /// the key as-is" separately from "emit a function of the key". + pub(crate) fn is_monotonic_function_of( + &self, + expr: &Arc, + range_key: &Arc, + ) -> bool { + if expr.eq(range_key) { + return false; + } + let dependencies = Dependencies::new(std::iter::once(PhysicalSortExpr::new( + Arc::clone(range_key), + Default::default(), + ))); + matches!( + get_expr_properties(expr, &dependencies, &self.schema) + .map(|properties| properties.sort_properties), + Ok(SortProperties::Ordered(_)) + ) + } + /// Transforms this `EquivalenceProperties` by mapping columns in the /// original schema to columns in the new schema by index. pub fn with_new_schema(mut self, schema: SchemaRef) -> Result { diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index b493a9b10d84e..277f11712cbe7 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -17,23 +17,20 @@ //! [`Partitioning`] and [`Distribution`] for `ExecutionPlans` -use crate::expressions::{Column, Literal, UnKnownColumn}; -use crate::utils::collect_columns; +use crate::expressions::{Literal, UnKnownColumn}; +use crate::simplifier::const_evaluator::create_dummy_batch; use crate::{ EquivalenceProperties, PhysicalExpr, equivalence::ProjectionMapping, physical_exprs_contains, physical_exprs_equal, }; pub use datafusion_common::SplitPoint; +use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{Result, ScalarValue, validate_range_split_points}; use datafusion_expr::ColumnarValue; -use datafusion_expr::interval_arithmetic::Interval; -use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; +use datafusion_expr::interval_arithmetic::checked_predecessor; use datafusion_physical_expr_common::physical_expr::format_physical_expr_list; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; -use arrow::array::new_null_array; -use arrow::datatypes::Schema; -use arrow::record_batch::RecordBatch; #[cfg(feature = "proto")] use datafusion_physical_expr_common::sort_expr::{ sort_exprs_try_from_proto, sort_exprs_try_to_proto, @@ -262,10 +259,9 @@ impl RangePartitioning { /// Returns `None` if any range key cannot be projected or if projection /// collapses distinct range keys into duplicate output expressions. /// - /// A range key that is not emitted as-is can still be projected when the - /// mapping contains a monotonic function of that key (for example - /// `date_bin(interval, timestamp)` or `date_trunc(unit, timestamp)` while - /// range-partitioned on `timestamp`). + /// If a projection drops a range key but keeps a monotonic function of it + /// (for example `date_bin(interval, timestamp)` or `date_trunc(unit, timestamp)` + /// while range-partitioned on `timestamp`), the range can still be projected. /// Adjacent partitions stay disjoint only when evaluating the function at /// each split point and its predecessor yields different values, so bins /// do not straddle file groups. @@ -291,7 +287,6 @@ impl RangePartitioning { &sort_expr.expr, &split_points, key_idx, - input_eq_properties.schema(), ) { return None; } @@ -300,7 +295,6 @@ impl RangePartitioning { &sort_expr.expr, &split_points, key_idx, - input_eq_properties.schema(), ) { split_points = updated; } @@ -325,67 +319,12 @@ fn monotonic_range_key_projection( eq_properties: &EquivalenceProperties, ) -> Option<(Arc, Arc)> { mapping.iter().find_map(|(source, targets)| { - is_order_preserving_function_of(source, &sort_expr.expr, eq_properties.schema()) + eq_properties + .is_monotonic_function_of(source, &sort_expr.expr) .then(|| (Arc::clone(&targets.first().0), Arc::clone(source))) }) } -/// Returns true when `expr` is a (possibly non-strict) monotonic function of -/// `range_key` plus literals, such as `date_bin(interval, timestamp)` or -/// `date_trunc(unit, timestamp)`. -fn is_order_preserving_function_of( - expr: &Arc, - range_key: &Arc, - schema: &arrow::datatypes::SchemaRef, -) -> bool { - if expr.eq(range_key) { - return false; - } - let expr_cols = collect_columns(expr); - let key_cols = collect_columns(range_key); - if key_cols.is_empty() || expr_cols != key_cols { - return false; - } - let Ok(child_properties) = expr - .children() - .iter() - .map(|child| function_child_properties(child, range_key, schema)) - .collect::>>() - else { - return false; - }; - matches!( - expr.get_properties(&child_properties) - .map(|properties| properties.sort_properties), - Ok(SortProperties::Ordered(_)) - ) -} - -fn function_child_properties( - child: &Arc, - range_key: &Arc, - schema: &arrow::datatypes::SchemaRef, -) -> Result { - if child.eq(range_key) { - let data_type = child.data_type(schema)?; - return Ok(ExprProperties { - sort_properties: SortProperties::Ordered(Default::default()), - range: Interval::make_unbounded(&data_type)?, - preserves_lex_ordering: true, - strictly_order_preserving: true, - }); - } - if child.downcast_ref::().is_some() { - return Ok(ExprProperties { - sort_properties: SortProperties::Singleton, - range: Interval::make_unbounded(&child.data_type(schema)?)?, - preserves_lex_ordering: true, - strictly_order_preserving: true, - }); - } - Ok(ExprProperties::new_unknown()) -} - /// Adjacent range partitions remain disjoint on `fn_expr` when the function /// value at each split differs from the value immediately below the split. fn monotonic_fn_keeps_partitions_disjoint( @@ -393,22 +332,18 @@ fn monotonic_fn_keeps_partitions_disjoint( range_key: &Arc, split_points: &[SplitPoint], key_idx: usize, - schema: &arrow::datatypes::SchemaRef, ) -> bool { split_points.iter().all(|split_point| { let Some(split_value) = split_point.values().get(key_idx) else { return false; }; - let Some(predecessor) = scalar_predecessor(split_value) else { + let Some(predecessor) = checked_predecessor(split_value) else { return false; }; - let Some(at_split) = - evaluate_expr_on_key(fn_expr, range_key, split_value, schema) - else { + let Some(at_split) = evaluate_expr_on_key(fn_expr, range_key, split_value) else { return false; }; - let Some(below_split) = - evaluate_expr_on_key(fn_expr, range_key, &predecessor, schema) + let Some(below_split) = evaluate_expr_on_key(fn_expr, range_key, &predecessor) else { return false; }; @@ -421,14 +356,12 @@ fn project_split_points_through_fn( range_key: &Arc, split_points: &[SplitPoint], key_idx: usize, - schema: &arrow::datatypes::SchemaRef, ) -> Option> { split_points .iter() .map(|split_point| { let split_value = split_point.values().get(key_idx)?; - let projected = - evaluate_expr_on_key(fn_expr, range_key, split_value, schema)?; + let projected = evaluate_expr_on_key(fn_expr, range_key, split_value)?; let mut values = split_point.values().to_vec(); values[key_idx] = projected; Some(SplitPoint::new(values)) @@ -436,62 +369,32 @@ fn project_split_points_through_fn( .collect() } +/// Evaluates `expr` after substituting `range_key` with `value`. fn evaluate_expr_on_key( expr: &Arc, range_key: &Arc, value: &ScalarValue, - schema: &Schema, ) -> Option { - let column = range_key.downcast_ref::()?; - // The table schema may mark columns non-nullable. Build a 1-row batch with - // nullable fields so unused columns can be null while still evaluating `expr`. - let nullable_schema = Arc::new(Schema::new( - schema - .fields() - .iter() - .map(|field| field.as_ref().clone().with_nullable(true)) - .collect::>(), - )); - let arrays = nullable_schema - .fields() - .iter() - .enumerate() - .map(|(idx, field)| { - if idx == column.index() { - value.to_array_of_size(1).ok() + let literal: Arc = Arc::new(Literal::new(value.clone())); + let rewritten = Arc::clone(expr) + .transform(|node| { + if node.eq(range_key) { + Ok(Transformed::yes(Arc::clone(&literal))) } else { - Some(new_null_array(field.data_type(), 1)) + Ok(Transformed::no(node)) } }) - .collect::>>()?; - let batch = RecordBatch::try_new(nullable_schema, arrays).ok()?; - match expr.evaluate(&batch).ok()? { + .ok()?; + if !rewritten.transformed { + return None; + } + let batch = create_dummy_batch().ok()?; + match rewritten.data.evaluate(batch).ok()? { ColumnarValue::Scalar(scalar) => Some(scalar), ColumnarValue::Array(array) => ScalarValue::try_from_array(&array, 0).ok(), } } -fn scalar_predecessor(value: &ScalarValue) -> Option { - match value { - ScalarValue::TimestampNanosecond(Some(v), tz) => Some( - ScalarValue::TimestampNanosecond(Some(v.checked_sub(1)?), tz.clone()), - ), - ScalarValue::TimestampMicrosecond(Some(v), tz) => Some( - ScalarValue::TimestampMicrosecond(Some(v.checked_sub(1)?), tz.clone()), - ), - ScalarValue::TimestampMillisecond(Some(v), tz) => Some( - ScalarValue::TimestampMillisecond(Some(v.checked_sub(1)?), tz.clone()), - ), - ScalarValue::TimestampSecond(Some(v), tz) => Some(ScalarValue::TimestampSecond( - Some(v.checked_sub(1)?), - tz.clone(), - )), - ScalarValue::Int64(Some(v)) => Some(ScalarValue::Int64(Some(v.checked_sub(1)?))), - ScalarValue::Int32(Some(v)) => Some(ScalarValue::Int32(Some(v.checked_sub(1)?))), - _ => None, - } -} - /// Range([x]) satisfies grouping by `(..., f(x), ...)` when `f` is monotonic in /// `x` and adjacent partitions do not share `f` values (bins do not straddle /// split points). That makes `(key, date_bin(timestamp))` and @@ -506,15 +409,13 @@ fn range_monotonic_fn_satisfies_keys( return false; } let range_key = &range.ordering()[0].expr; - let schema = eq_properties.schema(); required_exprs.iter().any(|required| { - is_order_preserving_function_of(required, range_key, schema) + eq_properties.is_monotonic_function_of(required, range_key) && monotonic_fn_keeps_partitions_disjoint( required, range_key, range.split_points(), 0, - schema, ) }) } From 693cd3ee58588dd519ef78bab96990a4ce3b1621 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Thu, 20 Aug 2026 14:47:36 -0400 Subject: [PATCH 6/8] fix: wrap Range([x]) in backticks to satisfy rustdoc link checks Co-authored-by: Cursor --- datafusion/physical-expr/src/partitioning.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 277f11712cbe7..e00e7f908e57c 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -395,7 +395,7 @@ fn evaluate_expr_on_key( } } -/// Range([x]) satisfies grouping by `(..., f(x), ...)` when `f` is monotonic in +/// `Range([x])` satisfies grouping by `(..., f(x), ...)` when `f` is monotonic in /// `x` and adjacent partitions do not share `f` values (bins do not straddle /// split points). That makes `(key, date_bin(timestamp))` and /// `(key, date_trunc(timestamp))` partition-disjoint when the table is From c36301c84833df66957fa6ca1a1e4f13a89a8074 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Thu, 20 Aug 2026 15:13:06 -0400 Subject: [PATCH 7/8] test: cover range date_bin/date_trunc edge cases and checked_predecessor Co-authored-by: Cursor --- .../expr-common/src/interval_arithmetic.rs | 26 ++++- datafusion/physical-expr/src/partitioning.rs | 105 ++++++++++++++++++ 2 files changed, 130 insertions(+), 1 deletion(-) diff --git a/datafusion/expr-common/src/interval_arithmetic.rs b/datafusion/expr-common/src/interval_arithmetic.rs index 3f476bb3791a5..bbd85f4971d8a 100644 --- a/datafusion/expr-common/src/interval_arithmetic.rs +++ b/datafusion/expr-common/src/interval_arithmetic.rs @@ -2275,7 +2275,8 @@ impl NullableInterval { mod tests { use crate::{ interval_arithmetic::{ - Interval, handle_overflow, next_value, prev_value, satisfy_greater, + Interval, checked_predecessor, handle_overflow, next_value, prev_value, + satisfy_greater, }, operator::Operator, }; @@ -2372,6 +2373,29 @@ mod tests { Ok(()) } + #[test] + fn test_checked_predecessor() { + assert_eq!( + checked_predecessor(&ScalarValue::Int64(Some(10))), + Some(ScalarValue::Int64(Some(9))) + ); + assert_eq!(checked_predecessor(&ScalarValue::Int64(None)), None); + assert_eq!( + checked_predecessor(&ScalarValue::Int64(Some(i64::MIN))), + None + ); + assert_eq!( + checked_predecessor(&ScalarValue::TimestampNanosecond(Some(i64::MIN), None)), + None + ); + // Types without a discrete predecessor return the same value from + // `prev_value`, which `checked_predecessor` treats as absent. + assert_eq!( + checked_predecessor(&ScalarValue::Utf8(Some("a".into()))), + None + ); + } + #[test] fn test_new_interval() -> Result<()> { use ScalarValue::*; diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index e00e7f908e57c..bfeccc39ea1a9 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -1489,6 +1489,32 @@ mod tests { PartitioningSatisfaction::NotSatisfied, ); + let min_ts = fixture.range_partitioning([1], vec![ts_ns_split(i64::MIN)]); + assert_satisfaction( + "type-minimum split has no predecessor so date_bin grouping is not disjoint", + &min_ts, + &required, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + let compound = fixture.range_partitioning( + [0, 1], + vec![SplitPoint::new(vec![ + ScalarValue::Utf8(Some("k".into())), + ScalarValue::TimestampNanosecond(Some(hour_ns), None), + ])], + ); + assert_satisfaction( + "multi-key Range([key, timestamp]) does not use single-key date_bin subset logic", + &compound, + &required, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + Ok(()) } @@ -1521,6 +1547,85 @@ mod tests { }; assert_eq!(partition_count, 2); + let min_ts = fixture.range_partitioning([0], vec![ts_ns_split(i64::MIN)]); + let projected = min_ts.project(&mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); + + Ok(()) + } + + #[test] + fn test_range_partitioning_project_through_date_trunc() -> Result<()> { + let fixture = PartitioningTestFixture::new(vec![( + "timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, None), + )])?; + let hour_ns = 1_704_070_800_000_000_000i64; + let trunc_hour = date_trunc_of(fixture.col(0), "hour"); + let target: Arc = Arc::new(Column::new("time_bin", 0)); + let mapping = ProjectionMapping::from_iter([( + Arc::clone(&trunc_hour), + ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), + )]); + + let aligned = fixture.range_partitioning([0], vec![ts_ns_split(hour_ns)]); + let projected = aligned.project(&mapping, &fixture.eq_properties); + assert_eq!( + projected.to_string(), + "Range([time_bin@0 ASC], [(1704070800000000000)], 2)" + ); + + let trunc_day = date_trunc_of(fixture.col(0), "day"); + let day_mapping = ProjectionMapping::from_iter([( + Arc::clone(&trunc_day), + ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), + )]); + let projected = aligned.project(&day_mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); + + Ok(()) + } + + #[test] + fn test_range_partitioning_project_compound_through_date_bin() -> Result<()> { + let fixture = PartitioningTestFixture::new(vec![ + ("key", DataType::Utf8), + ("timestamp", DataType::Timestamp(TimeUnit::Nanosecond, None)), + ])?; + let hour_ns = 1_704_070_800_000_000_000i64; + let date_bin = date_bin_of(fixture.col(1), 60_000_000_000); + let key_target: Arc = Arc::new(Column::new("key", 0)); + let bin_target: Arc = Arc::new(Column::new("time_bin", 1)); + let mapping = ProjectionMapping::from_iter([ + ( + fixture.col(0), + ProjectionTargets::from(vec![(Arc::clone(&key_target), 0)]), + ), + ( + Arc::clone(&date_bin), + ProjectionTargets::from(vec![(Arc::clone(&bin_target), 1)]), + ), + ]); + + let aligned = fixture.range_partitioning( + [0, 1], + vec![SplitPoint::new(vec![ + ScalarValue::Utf8(Some("k".into())), + ScalarValue::TimestampNanosecond(Some(hour_ns), None), + ])], + ); + let projected = aligned.project(&mapping, &fixture.eq_properties); + assert_eq!( + projected.to_string(), + "Range([key@0 ASC, time_bin@1 ASC], [(k, 1704070800000000000)], 2)" + ); + Ok(()) } From 0c118e0785c4925e66af475399d9ddad01e0cc15 Mon Sep 17 00:00:00 2001 From: Nga Tran Date: Thu, 20 Aug 2026 15:59:21 -0400 Subject: [PATCH 8/8] test: fail closed when date_bin range splits are null or invalid Co-authored-by: Cursor --- datafusion/physical-expr/src/partitioning.rs | 87 ++++++++++++++++++++ 1 file changed, 87 insertions(+) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index bfeccc39ea1a9..b87eb60cf8340 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -1515,6 +1515,61 @@ mod tests { PartitioningSatisfaction::NotSatisfied, ); + let bin_only = Distribution::KeyPartitioned(vec![date_bin_of( + fixture.col(1), + 60_000_000_000, + )]); + assert_satisfaction( + "aligned hour split: Range(timestamp) subset-satisfies GROUP BY date_bin(60s, timestamp)", + &aligned, + &bin_only, + &fixture.eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + + let null_split = fixture.range_partitioning( + [1], + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + None, None, + )])], + ); + assert_satisfaction( + "null split has no predecessor so date_bin grouping is not disjoint", + &null_split, + &required, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + // `RangePartitioning::new` skips validation, so disjointness must still + // fail closed on split points that do not match the range key. + let empty_split = Partitioning::Range(RangePartitioning::new( + fixture.range_ordering([1]), + vec![SplitPoint::new(vec![])], + )); + assert_satisfaction( + "split point missing the range key is not disjoint for date_bin grouping", + &empty_split, + &required, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + let mismatched_split = Partitioning::Range(RangePartitioning::new( + fixture.range_ordering([1]), + vec![int_split_point([10])], + )); + assert_satisfaction( + "non-timestamp split cannot be evaluated as date_bin, so grouping is not disjoint", + &mismatched_split, + &required, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + Ok(()) } @@ -1554,6 +1609,38 @@ mod tests { }; assert_eq!(partition_count, 2); + let null_split = fixture.range_partitioning( + [0], + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + None, None, + )])], + ); + let projected = null_split.project(&mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); + + let empty_split = Partitioning::Range(RangePartitioning::new( + fixture.range_ordering([0]), + vec![SplitPoint::new(vec![])], + )); + let projected = empty_split.project(&mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); + + let mismatched_split = Partitioning::Range(RangePartitioning::new( + fixture.range_ordering([0]), + vec![int_split_point([10])], + )); + let projected = mismatched_split.project(&mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); + Ok(()) }