diff --git a/datafusion/expr-common/src/interval_arithmetic.rs b/datafusion/expr-common/src/interval_arithmetic.rs index 68541e1e6b32c..bbd85f4971d8a 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; } @@ -2261,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, }; @@ -2358,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/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index f52b320ed284f..75d3c8a0b97ef 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -1319,6 +1319,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 98f082f7256db..b87eb60cf8340 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -17,14 +17,20 @@ //! [`Partitioning`] and [`Distribution`] for `ExecutionPlans` +use crate::expressions::{Literal, UnKnownColumn}; +use crate::simplifier::const_evaluator::create_dummy_batch; 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::tree_node::{Transformed, TreeNode}; +use datafusion_common::{Result, ScalarValue, validate_range_split_points}; +use datafusion_expr::ColumnarValue; +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}; + #[cfg(feature = "proto")] use datafusion_physical_expr_common::sort_expr::{ sort_exprs_try_from_proto, sort_exprs_try_to_proto, @@ -252,25 +258,48 @@ impl RangePartitioning { /// /// Returns `None` if any range key cannot be projected or if projection /// collapses distinct range keys into duplicate output expressions. + /// + /// 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. 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, + ) { + return None; + } + if let Some(updated) = project_split_points_through_fn( + &source, + &sort_expr.expr, + &split_points, + key_idx, + ) { + 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,11 +307,119 @@ 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)| { + eq_properties + .is_monotonic_function_of(source, &sort_expr.expr) + .then(|| (Arc::clone(&targets.first().0), Arc::clone(source))) + }) +} + +/// 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, +) -> bool { + split_points.iter().all(|split_point| { + let Some(split_value) = split_point.values().get(key_idx) else { + return false; + }; + let Some(predecessor) = checked_predecessor(split_value) else { + return false; + }; + 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) + 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, +) -> 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)?; + let mut values = split_point.values().to_vec(); + values[key_idx] = projected; + Some(SplitPoint::new(values)) + }) + .collect() +} + +/// Evaluates `expr` after substituting `range_key` with `value`. +fn evaluate_expr_on_key( + expr: &Arc, + range_key: &Arc, + value: &ScalarValue, +) -> Option { + 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 { + Ok(Transformed::no(node)) + } + }) + .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(), + } +} + +/// `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; + required_exprs.iter().any(|required| { + eq_properties.is_monotonic_function_of(required, range_key) + && monotonic_fn_keeps_partitions_disjoint( + required, + range_key, + range.split_points(), + 0, + ) + }) +} + impl Display for RangePartitioning { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { let split_points = format_range_split_points(&self.split_points); @@ -433,12 +570,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 +895,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 +1378,344 @@ 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, + ); + + 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, + ); + + 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(()) + } + + #[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); + + 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); + + 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(()) + } + + #[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(()) + } + #[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 ##########