From 6073032cddcddf7ee8d6582152e64ac30a66123d Mon Sep 17 00:00:00 2001 From: "ben.bellick" Date: Wed, 16 Sep 2026 18:20:06 +0000 Subject: [PATCH 1/4] feat: add audited range analysis capability --- datafusion/expr/src/udf.rs | 70 +++++++ .../functions/src/datetime/date_trunc.rs | 187 +++++++++++++++++- .../physical-expr/src/scalar_function.rs | 50 +++++ 3 files changed, 305 insertions(+), 2 deletions(-) diff --git a/datafusion/expr/src/udf.rs b/datafusion/expr/src/udf.rs index 4340fbb4a384f..d3ac5fe3e5888 100644 --- a/datafusion/expr/src/udf.rs +++ b/datafusion/expr/src/udf.rs @@ -364,6 +364,15 @@ impl ScalarUDF { self.inner.output_ordering(inputs) } + /// See [`ScalarUDFImpl::supports_range_partitioning_analysis`] for more details. + pub fn supports_range_partitioning_analysis( + &self, + argument_types: &[DataType], + ) -> bool { + self.inner + .supports_range_partitioning_analysis(argument_types) + } + pub fn preserves_lex_ordering(&self, inputs: &[ExprProperties]) -> Result { self.inner.preserves_lex_ordering(inputs) } @@ -967,6 +976,38 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { Ok(inner(inputs)) } + /// Returns whether this function and argument-type combination may + /// participate in range-partitioning transformation analysis. + /// + /// Ordinary [`SortProperties::Ordered`] metadata is not sufficient for this + /// analysis. A non-strictly ordered function can map values from opposite + /// sides of a range split to the same output, and successful evaluation of + /// some functions can produce null from non-null input. Either behavior can + /// place one transformed equality group in multiple partitions. + /// + /// Returning `true` asserts that this argument-type combination has been + /// audited with the following contract: + /// + /// - whenever the concrete expression reports same-direction ordering with + /// all non-range inputs singleton, it is globally nondecreasing over the + /// admitted non-null input domain; + /// - successful evaluation preserves nullness: null range inputs remain + /// null and non-null range inputs do not produce null; and + /// - successful scalar evaluation at a split boundary is consistent with + /// array evaluation of that value. + /// + /// This method does not assert that a concrete range partitioning is + /// preserved. Callers must still identify the range input, establish + /// same-direction ordering and singleton arguments, verify exact split + /// types, and prove at every split that the images of adjacent partitions + /// do not overlap. Evaluation errors must fail closed. + /// + /// The default is `false`. This capability is not currently transported + /// across FFI, so foreign UDFs also remain denied by default. + fn supports_range_partitioning_analysis(&self, _argument_types: &[DataType]) -> bool { + false + } + /// Returns true if the function preserves lexicographical ordering based on /// the input ordering. /// @@ -1189,6 +1230,11 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { self.inner.output_ordering(inputs) } + fn supports_range_partitioning_analysis(&self, argument_types: &[DataType]) -> bool { + self.inner + .supports_range_partitioning_analysis(argument_types) + } + fn preserves_lex_ordering(&self, inputs: &[ExprProperties]) -> Result { self.inner.preserves_lex_ordering(inputs) } @@ -1221,6 +1267,7 @@ mod tests { name: &'static str, field: &'static str, signature: Signature, + supports_range_partitioning_analysis: bool, } impl ScalarUDFImpl for TestScalarUDFImpl { fn name(&self) -> &str { @@ -1238,6 +1285,14 @@ mod tests { fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { unimplemented!() } + + fn supports_range_partitioning_analysis( + &self, + argument_types: &[DataType], + ) -> bool { + self.supports_range_partitioning_analysis + && argument_types == [DataType::Int64] + } } // PartialEq and Hash must be consistent, and also PartialEq and PartialOrd @@ -1271,11 +1326,26 @@ mod tests { assert_eq!(b.partial_cmp(&o), Some(Ordering::Less)); } + #[test] + fn aliases_preserve_range_partitioning_analysis_support() { + let function = ScalarUDF::from(TestScalarUDFImpl { + name: "range_function", + field: "a", + signature: Signature::any(1, Volatility::Immutable), + supports_range_partitioning_analysis: true, + }) + .with_aliases(["range_function_alias"]); + + assert!(function.supports_range_partitioning_analysis(&[DataType::Int64])); + assert!(!function.supports_range_partitioning_analysis(&[DataType::Utf8])); + } + fn test_func(name: &'static str, parameter: &'static str) -> ScalarUDF { ScalarUDF::from(TestScalarUDFImpl { name, field: parameter, signature: Signature::any(1, Volatility::Immutable), + supports_range_partitioning_analysis: false, }) } diff --git a/datafusion/functions/src/datetime/date_trunc.rs b/datafusion/functions/src/datetime/date_trunc.rs index e61637c0931db..39ecaa06ad782 100644 --- a/datafusion/functions/src/datetime/date_trunc.rs +++ b/datafusion/functions/src/datetime/date_trunc.rs @@ -441,6 +441,19 @@ impl ScalarUDFImpl for DateTruncFunc { Ok(SortProperties::Unordered) } } + + fn supports_range_partitioning_analysis(&self, argument_types: &[DataType]) -> bool { + // Named timezone transitions can make date_trunc non-monotonic (#25353). + // Keep the audited domain to timezone-less timestamp columns. + matches!( + argument_types, + [ + DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8, + Timestamp(_, None) + ] + ) + } + fn documentation(&self) -> Option<&Documentation> { self.doc() } @@ -883,7 +896,10 @@ mod tests { }; use arrow::array::cast::as_primitive_array; - use arrow::array::types::{ArrowTimestampType, TimestampNanosecondType}; + use arrow::array::types::{ + ArrowTimestampType, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, + }; use arrow::array::{ Array, PrimitiveArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, @@ -891,11 +907,178 @@ mod tests { use arrow::buffer::NullBuffer; use arrow::compute::DatePart; use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos; - use arrow::datatypes::{DataType, Field, TimeUnit}; + use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit}; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; + #[test] + fn supports_range_partitioning_analysis_for_timezone_less_timestamps() { + let function = DateTruncFunc::new(); + + for unit in [ + TimeUnit::Second, + TimeUnit::Millisecond, + TimeUnit::Microsecond, + TimeUnit::Nanosecond, + ] { + for string_type in [DataType::Utf8, DataType::Utf8View, DataType::LargeUtf8] { + assert!(function.supports_range_partitioning_analysis(&[ + string_type, + DataType::Timestamp(unit, None), + ])); + } + + for timezone in ["UTC", "+05:30", "America/Goose_Bay"] { + assert!(!function.supports_range_partitioning_analysis(&[ + DataType::Utf8, + DataType::Timestamp(unit, Some(timezone.into())), + ])); + } + } + + assert!(!function.supports_range_partitioning_analysis(&[ + DataType::Utf8, + DataType::Time32(TimeUnit::Second), + ])); + assert!(!function.supports_range_partitioning_analysis(&[ + DataType::Utf8, + DataType::Time64(TimeUnit::Nanosecond), + ])); + assert!(!function.supports_range_partitioning_analysis(&[DataType::Utf8,])); + assert!(!function.supports_range_partitioning_analysis(&[ + DataType::Int64, + DataType::Timestamp(TimeUnit::Nanosecond, None), + ])); + } + + #[test] + fn date_bin_does_not_support_range_partitioning_analysis() { + assert!( + !crate::datetime::date_bin::DateBinFunc::new() + .supports_range_partitioning_analysis(&[ + DataType::Interval(IntervalUnit::MonthDayNano), + DataType::Timestamp(TimeUnit::Nanosecond, None), + ]) + ); + } + + fn invoke_date_trunc( + granularity: &str, + timestamp: ColumnarValue, + timestamp_type: DataType, + number_rows: usize, + ) -> ColumnarValue { + DateTruncFunc::new() + .invoke_with_args(ScalarFunctionArgs { + args: vec![ScalarValue::from(granularity).into(), timestamp], + arg_fields: vec![ + Field::new("granularity", DataType::Utf8, false).into(), + Field::new("timestamp", timestamp_type.clone(), true).into(), + ], + number_rows, + return_field: Field::new("date_trunc", timestamp_type, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap() + } + + fn assert_timezone_less_range_analysis_contract() { + let scale = match T::UNIT { + TimeUnit::Second => 1_000_000_000, + TimeUnit::Millisecond => 1_000_000, + TimeUnit::Microsecond => 1_000, + TimeUnit::Nanosecond => 1, + }; + let values = [ + Some("1969-12-28T23:59:59Z"), + Some("1969-12-31T23:59:59Z"), + Some("1970-01-01T00:00:00Z"), + Some("1999-12-31T23:59:59Z"), + Some("2000-01-01T00:00:00Z"), + Some("2000-02-28T23:59:59Z"), + Some("2000-02-29T00:00:00Z"), + Some("2000-03-01T00:00:00Z"), + Some("2023-12-31T23:59:59Z"), + Some("2024-01-01T00:00:00Z"), + None, + ] + .map(|value| { + value.map(|value| string_to_timestamp_nanos(value).unwrap() / scale) + }); + let timestamp_type = DataType::Timestamp(T::UNIT, None); + + for granularity in [ + "microsecond", + "millisecond", + "second", + "minute", + "hour", + "day", + "week", + "month", + "quarter", + "year", + ] { + let input = PrimitiveArray::::from_iter(values); + let output = invoke_date_trunc( + granularity, + ColumnarValue::Array(Arc::new(input.clone())), + timestamp_type.clone(), + input.len(), + ); + let ColumnarValue::Array(output) = output else { + panic!("expected array output") + }; + let output = as_primitive_array::(&output); + + assert_eq!( + input.nulls(), + output.nulls(), + "{granularity} must preserve nullness for {:?}", + T::UNIT + ); + + let mut previous = None; + for value in output.iter().flatten() { + if let Some(previous) = previous { + assert!( + previous <= value, + "{granularity} must be nondecreasing for {:?}", + T::UNIT + ); + } + previous = Some(value); + } + + for (index, value) in values.iter().enumerate() { + let scalar_output = invoke_date_trunc( + granularity, + ColumnarValue::Scalar(ScalarValue::new_timestamp::(*value, None)), + timestamp_type.clone(), + 1, + ); + let ColumnarValue::Scalar(scalar_output) = scalar_output else { + panic!("expected scalar output") + }; + assert_eq!( + scalar_output, + ScalarValue::try_from_array(output, index).unwrap(), + "scalar and array evaluation must agree for {granularity} {:?}", + T::UNIT + ); + } + } + } + + #[test] + fn timezone_less_timestamps_satisfy_range_analysis_contract() { + assert_timezone_less_range_analysis_contract::(); + assert_timezone_less_range_analysis_contract::(); + assert_timezone_less_range_analysis_contract::(); + assert_timezone_less_range_analysis_contract::(); + } + #[test] fn date_trunc_test() { let cases = vec![ diff --git a/datafusion/physical-expr/src/scalar_function.rs b/datafusion/physical-expr/src/scalar_function.rs index 6a5ab219aa8dd..91b8406bf4e1a 100644 --- a/datafusion/physical-expr/src/scalar_function.rs +++ b/datafusion/physical-expr/src/scalar_function.rs @@ -158,6 +158,22 @@ impl ScalarFunctionExpr { &self.config_options } + /// Returns whether this function expression's argument types have been + /// explicitly audited as candidates for range-partitioning analysis. + pub fn supports_range_partitioning_analysis( + &self, + input_schema: &Schema, + ) -> Result { + let argument_types = self + .args + .iter() + .map(|argument| argument.data_type(input_schema)) + .collect::>>()?; + Ok(self + .fun + .supports_range_partitioning_analysis(&argument_types)) + } + /// Given an arbitrary PhysicalExpr attempt to downcast it to a ScalarFunctionExpr /// and verify that its inner function is of type T. /// If the downcast fails, or the function is not of type T, returns `None`. @@ -383,6 +399,40 @@ mod tests { fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(42)))) } + + fn supports_range_partitioning_analysis( + &self, + argument_types: &[DataType], + ) -> bool { + argument_types == [DataType::Float32] + } + } + + #[test] + fn test_range_partitioning_analysis_support_uses_argument_types() { + let function = Arc::new(ScalarUDF::from(MockScalarUDF { + signature: Signature::any(1, Volatility::Immutable), + })); + let config_options = Arc::new(ConfigOptions::new()); + + for (data_type, expected) in [(DataType::Float32, true), (DataType::Int32, false)] + { + let schema = Schema::new(vec![Field::new("a", data_type, false)]); + let expression = ScalarFunctionExpr::try_new( + Arc::clone(&function), + vec![Arc::new(Column::new("a", 0))], + &schema, + Arc::clone(&config_options), + ) + .unwrap(); + + assert_eq!( + expression + .supports_range_partitioning_analysis(&schema) + .unwrap(), + expected + ); + } } #[test] From c88e7e21507f771e76985ec6f447852011a6c178 Mon Sep 17 00:00:00 2001 From: "ben.bellick" Date: Wed, 16 Sep 2026 22:05:31 +0000 Subject: [PATCH 2/4] feat: use range transformations for aggregation --- .../expr-common/src/interval_arithmetic.rs | 61 ++- .../src/equivalence/properties/mod.rs | 24 + datafusion/physical-expr/src/partitioning.rs | 501 +++++++++++++++++- .../enforce_distribution.rs | 67 ++- .../physical-optimizer/src/sanity_checker.rs | 22 +- .../test_files/range_sorted_time_bin_agg.slt | 58 ++ 6 files changed, 704 insertions(+), 29 deletions(-) diff --git a/datafusion/expr-common/src/interval_arithmetic.rs b/datafusion/expr-common/src/interval_arithmetic.rs index 9f1291353dc29..3e4d37c56b11c 100644 --- a/datafusion/expr-common/src/interval_arithmetic.rs +++ b/datafusion/expr-common/src/interval_arithmetic.rs @@ -1243,6 +1243,34 @@ fn prev_value(value: ScalarValue) -> ScalarValue { value_transition!(MIN, false, value) } +/// Returns the previous distinct value, or `None` for nulls, non-finite +/// floats, type minima, and types without a discrete predecessor. +pub fn checked_predecessor(value: &ScalarValue) -> Option { + checked_adjacent(value, false) +} + +/// Returns the next distinct value, or `None` for nulls, non-finite floats, +/// type maxima, and types without a discrete successor. +pub fn checked_successor(value: &ScalarValue) -> Option { + checked_adjacent(value, true) +} + +fn checked_adjacent(value: &ScalarValue, successor: bool) -> Option { + if value.is_null() + || matches!(value, ScalarValue::Float32(Some(v)) if !v.is_finite()) + || matches!(value, ScalarValue::Float64(Some(v)) if !v.is_finite()) + { + return None; + } + + let adjacent = if successor { + next_value(value.clone()) + } else { + prev_value(value.clone()) + }; + (!adjacent.is_null() && adjacent != *value).then_some(adjacent) +} + trait OneTrait: Sized + std::ops::Add + std::ops::Sub { fn one() -> Self; } @@ -2261,7 +2289,8 @@ impl NullableInterval { mod tests { use crate::{ interval_arithmetic::{ - Interval, handle_overflow, next_value, prev_value, satisfy_greater, + Interval, checked_predecessor, checked_successor, handle_overflow, + next_value, prev_value, satisfy_greater, }, operator::Operator, }; @@ -2358,6 +2387,36 @@ mod tests { Ok(()) } + #[test] + fn test_checked_adjacent_values() { + assert_eq!( + checked_predecessor(&ScalarValue::Int64(Some(10))), + Some(ScalarValue::Int64(Some(9))) + ); + assert_eq!( + checked_successor(&ScalarValue::Int64(Some(10))), + Some(ScalarValue::Int64(Some(11))) + ); + assert_eq!(checked_predecessor(&ScalarValue::Int64(None)), None); + assert_eq!( + checked_predecessor(&ScalarValue::Int64(Some(i64::MIN))), + None + ); + assert_eq!(checked_successor(&ScalarValue::Int64(Some(i64::MAX))), None); + assert_eq!( + checked_predecessor(&ScalarValue::Utf8(Some("a".into()))), + None + ); + assert_eq!( + checked_successor(&ScalarValue::Float64(Some(f64::NAN))), + None + ); + assert_eq!( + checked_predecessor(&ScalarValue::Float32(Some(f32::NEG_INFINITY))), + None + ); + } + #[test] fn test_new_interval() -> Result<()> { use ScalarValue::*; diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index c68157fecbd8c..3d01495458de2 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -39,6 +39,7 @@ use crate::{ PhysicalSortRequirement, }; +use arrow::compute::SortOptions; use arrow::datatypes::SchemaRef; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{Constraint, Constraints, HashMap, Result, plan_err}; @@ -1407,6 +1408,29 @@ impl EquivalenceProperties { .unwrap_or_else(|_| ExprProperties::new_unknown()) } + /// Returns whether `transformed_expr` reports the same ordering direction + /// when `source_expr` is ordered and all other inputs are singleton. + pub(crate) fn is_same_direction_ordered_transform( + &self, + transformed_expr: &Arc, + source_expr: &Arc, + ) -> bool { + if transformed_expr.eq(source_expr) { + return false; + } + + let options = SortOptions::default(); + let dependencies = Dependencies::new(std::iter::once(PhysicalSortExpr::new( + Arc::clone(source_expr), + options, + ))); + matches!( + get_expr_properties(transformed_expr, &dependencies, &self.schema) + .map(|properties| properties.sort_properties), + Ok(SortProperties::Ordered(actual)) if actual == options + ) + } + /// 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 052fec6ebe97e..6ddcd4df82a47 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -17,13 +17,18 @@ //! [`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, + EquivalenceProperties, PhysicalExpr, ScalarFunctionExpr, + equivalence::ProjectionMapping, physical_exprs_contains, physical_exprs_equal, }; use arrow::datatypes::Schema; 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, checked_successor}; use datafusion_physical_expr_common::physical_expr::format_physical_expr_list; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; #[cfg(feature = "proto")] @@ -379,6 +384,170 @@ fn normalize_exprs( .collect() } +fn equivalent_expr( + left: &Arc, + right: &Arc, + eq_properties: &EquivalenceProperties, +) -> bool { + equivalent_exprs( + std::slice::from_ref(left), + std::slice::from_ref(right), + eq_properties, + ) +} + +fn evaluate_with_range_value( + 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; + } + + match rewritten.data.evaluate(create_dummy_batch().ok()?).ok()? { + ColumnarValue::Scalar(value) => Some(value), + ColumnarValue::Array(array) => ScalarValue::try_from_array(&array, 0).ok(), + } +} + +fn projected_split_points( + function: &Arc, + range_key: &Arc, + range: &RangePartitioning, + key_index: usize, + schema: &Schema, +) -> Option> { + let range_type = range_key.data_type(schema).ok()?; + range + .split_points() + .iter() + .map(|split_point| { + let split_value = split_point.values().get(key_index)?; + if split_value.data_type() != range_type { + return None; + } + let adjacent = if range.ordering()[key_index].options.descending { + checked_successor(split_value) + } else { + checked_predecessor(split_value) + }?; + let at_split = evaluate_with_range_value(function, range_key, split_value)?; + let across_split = evaluate_with_range_value(function, range_key, &adjacent)?; + if at_split.is_null() + || across_split.is_null() + || at_split.data_type() != across_split.data_type() + || at_split == across_split + { + return None; + } + + let mut values = split_point.values().to_vec(); + values[key_index] = at_split; + Some(SplitPoint::new(values)) + }) + .collect() +} + +fn is_audited_range_transform( + candidate: &Arc, + range_key: &Arc, + range: &RangePartitioning, + key_index: usize, + eq_properties: &EquivalenceProperties, +) -> bool { + let Some(function) = candidate.downcast_ref::() else { + return false; + }; + let schema = eq_properties.schema(); + if !function + .supports_range_partitioning_analysis(schema) + .unwrap_or(false) + || !eq_properties.is_same_direction_ordered_transform(candidate, range_key) + { + return false; + } + + let mut range_argument_count = 0; + let dummy_batch = create_dummy_batch().ok(); + for argument in function.args() { + if equivalent_expr(argument, range_key, eq_properties) { + range_argument_count += 1; + } else if dummy_batch + .as_ref() + .and_then(|batch| argument.evaluate(batch).ok()) + .is_none_or(|value| !matches!(value, ColumnarValue::Scalar(_))) + { + return false; + } + } + + range_argument_count == 1 + && projected_split_points(candidate, range_key, range, key_index, schema) + .is_some() +} + +fn range_transform_satisfaction( + range: &RangePartitioning, + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, + allow_subset: bool, +) -> PartitioningSatisfaction { + if required_exprs.is_empty() || range.ordering().len() > required_exprs.len() { + return PartitioningSatisfaction::NotSatisfied; + } + + let mut unmatched_required = (0..required_exprs.len()).collect::>(); + let last_range_key = range.ordering().len() - 1; + let mut used_transform = false; + + for (key_index, range_sort) in range.ordering().iter().enumerate() { + let mut match_position = unmatched_required.iter().position(|required_index| { + equivalent_expr( + &range_sort.expr, + &required_exprs[*required_index], + eq_properties, + ) + }); + if match_position.is_none() && key_index == last_range_key { + match_position = unmatched_required.iter().position(|required_index| { + is_audited_range_transform( + &required_exprs[*required_index], + &range_sort.expr, + range, + key_index, + eq_properties, + ) + }); + } + let Some(match_position) = match_position else { + return PartitioningSatisfaction::NotSatisfied; + }; + used_transform |= !equivalent_expr( + &range_sort.expr, + &required_exprs[unmatched_required[match_position]], + eq_properties, + ); + unmatched_required.swap_remove(match_position); + } + + if used_transform && (unmatched_required.is_empty() || allow_subset) { + PartitioningSatisfaction::KeyLocal + } else { + PartitioningSatisfaction::NotSatisfied + } +} + /// Represents how a [`Partitioning`] satisfies a [`Distribution`] requirement. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PartitioningSatisfaction { @@ -388,6 +557,12 @@ pub enum PartitioningSatisfaction { Exact, /// The partitioning satisfies the distribution requirement via subset logic Subset, + /// Equal values of the required keys are partition-local, but the concrete + /// partition layout is not claimed to match the requested distribution. + /// + /// This is sufficient for independent keyed aggregation. It is not general + /// distribution satisfaction and must not be used for co-partitioned joins. + KeyLocal, } impl PartitioningSatisfaction { @@ -395,6 +570,11 @@ impl PartitioningSatisfaction { matches!(self, Self::Exact | Self::Subset) } + /// Returns whether equal required-key tuples cannot occur in two partitions. + pub fn is_key_local(&self) -> bool { + !matches!(self, Self::NotSatisfied) + } + pub fn is_subset(&self) -> bool { *self == Self::Subset } @@ -437,8 +617,15 @@ impl Partitioning { == PartitioningSatisfaction::Exact } - /// Returns how this [`Partitioning`] satisfies the partitioning scheme mandated - /// by the `required` [`Distribution`]. + /// Returns the relationship between this [`Partitioning`] and the + /// `required` [`Distribution`]. + /// + /// [`PartitioningSatisfaction::KeyLocal`] is deliberately weaker than + /// distribution satisfaction: it only proves that equal required-key tuples + /// cannot occur in different partitions. Callers that require a concrete + /// partition layout must use [`PartitioningSatisfaction::is_satisfied`]. + /// Operators that only require independent processing of each key may use + /// [`PartitioningSatisfaction::is_key_local`]. #[expect( deprecated, reason = "HashPartitioned is accepted during the KeyPartitioned migration" @@ -477,12 +664,22 @@ 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.is_satisfied() { + satisfaction + } else { + range_transform_satisfaction( + range, + required_exprs, + eq_properties, + allow_subset, + ) + } } Partitioning::RoundRobinBatch(_) | Partitioning::UnknownPartitioning(_) => { @@ -814,7 +1011,8 @@ mod tests { 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 { @@ -946,6 +1144,293 @@ mod tests { ); } + fn date_trunc_of( + timestamp: Arc, + precision: &str, + timezone: Option<&str>, + ) -> Arc { + let timezone = timezone.map(Into::into); + Arc::new(ScalarFunctionExpr::new( + "date_trunc", + datafusion_functions::datetime::date_trunc(), + vec![ + Arc::new(Literal::new(ScalarValue::Utf8(Some(precision.into())))), + timestamp, + ], + Field::new( + "date_trunc", + DataType::Timestamp(TimeUnit::Nanosecond, timezone), + true, + ) + .into(), + Arc::new(ConfigOptions::default()), + )) + } + + fn date_bin_of(timestamp: Arc) -> Arc { + Arc::new(ScalarFunctionExpr::new( + "date_bin", + datafusion_functions::datetime::date_bin(), + vec![ + Arc::new(Literal::new(ScalarValue::new_interval_mdn( + 0, + 0, + 60_000_000_000, + ))), + timestamp, + ], + Field::new( + "date_bin", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + ) + .into(), + Arc::new(ConfigOptions::default()), + )) + } + + fn timestamp_split(value: Option, timezone: Option<&str>) -> SplitPoint { + SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + value, + timezone.map(Into::into), + )]) + } + + #[test] + fn range_transform_satisfaction_is_key_local_only() -> Result<()> { + let fixture = PartitioningTestFixture::new(vec![ + ("timestamp", DataType::Timestamp(TimeUnit::Nanosecond, None)), + ("tag", DataType::Utf8), + ])?; + let hour = 1_704_070_800_000_000_000; + let aligned = + fixture.range_partitioning([0], vec![timestamp_split(Some(hour), None)]); + let trunc_hour = date_trunc_of(fixture.col(0), "hour", None); + + let exact_keys = Distribution::KeyPartitioned(vec![Arc::clone(&trunc_hour)]); + let satisfaction = + aligned.satisfaction(&exact_keys, &fixture.eq_properties, false); + assert_eq!(satisfaction, PartitioningSatisfaction::KeyLocal); + assert!(satisfaction.is_key_local()); + assert!(!satisfaction.is_satisfied()); + + assert_satisfaction( + "additional grouping keys require subset permission", + &aligned, + &Distribution::KeyPartitioned(vec![trunc_hour, fixture.col(1)]), + &fixture.eq_properties, + PartitioningSatisfaction::KeyLocal, + PartitioningSatisfaction::NotSatisfied, + ); + + assert_satisfaction( + "date_bin remains default-denied", + &aligned, + &Distribution::KeyPartitioned(vec![date_bin_of(fixture.col(0))]), + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + let trunc_day = date_trunc_of(fixture.col(0), "day", None); + assert_satisfaction( + "an hour split straddles a day bucket", + &aligned, + &Distribution::KeyPartitioned(vec![trunc_day]), + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + let two_aligned = fixture.range_partitioning( + [0], + vec![ + timestamp_split(Some(hour), None), + timestamp_split(Some(hour + 3_600_000_000_000), None), + ], + ); + assert_satisfaction( + "every split must be aligned", + &two_aligned, + &exact_keys, + &fixture.eq_properties, + PartitioningSatisfaction::KeyLocal, + PartitioningSatisfaction::KeyLocal, + ); + + let unaligned = + fixture.range_partitioning([0], vec![timestamp_split(Some(hour + 1), None)]); + assert_satisfaction( + "an unaligned split straddles an hour bucket", + &unaligned, + &exact_keys, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + Ok(()) + } + + #[test] + fn range_transform_supports_only_the_final_compound_key() -> Result<()> { + let fixture = PartitioningTestFixture::new(vec![ + ("key", DataType::Utf8), + ("timestamp", DataType::Timestamp(TimeUnit::Nanosecond, None)), + ])?; + let hour = 1_704_070_800_000_000_000; + let split = SplitPoint::new(vec![ + ScalarValue::Utf8(Some("m".into())), + ScalarValue::TimestampNanosecond(Some(hour), None), + ]); + let range = fixture.range_partitioning([0, 1], vec![split]); + assert_satisfaction( + "the final compound key can be transformed", + &range, + &Distribution::KeyPartitioned(vec![ + fixture.col(0), + date_trunc_of(fixture.col(1), "hour", None), + ]), + &fixture.eq_properties, + PartitioningSatisfaction::KeyLocal, + PartitioningSatisfaction::KeyLocal, + ); + + let split = SplitPoint::new(vec![ + ScalarValue::TimestampNanosecond(Some(hour), None), + ScalarValue::Utf8(Some("m".into())), + ]); + let range = fixture.range_partitioning([1, 0], vec![split]); + assert_satisfaction( + "a non-final compound key cannot be transformed", + &range, + &Distribution::KeyPartitioned(vec![ + date_trunc_of(fixture.col(1), "hour", None), + fixture.col(0), + ]), + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + Ok(()) + } + + #[test] + fn range_transform_satisfaction_fails_closed() -> Result<()> { + let fixture = PartitioningTestFixture::new(vec![( + "timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, None), + )])?; + let trunc_hour = date_trunc_of(fixture.col(0), "hour", None); + let required = Distribution::KeyPartitioned(vec![trunc_hour]); + + for (description, split) in [ + ("null", timestamp_split(None, None)), + ("minimum", timestamp_split(Some(i64::MIN), None)), + ] { + assert_satisfaction( + description, + &fixture.range_partitioning([0], vec![split]), + &required, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + } + + let mismatched_split = Partitioning::Range(RangePartitioning::new( + fixture.range_ordering([0]), + vec![SplitPoint::new(vec![ScalarValue::Int64(Some(10))])], + )); + assert_satisfaction( + "split scalar type must match the range key", + &mismatched_split, + &required, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + let invalid_precision = Distribution::KeyPartitioned(vec![date_trunc_of( + fixture.col(0), + "not-a-precision", + None, + )]); + let aligned = fixture.range_partitioning( + [0], + vec![timestamp_split(Some(1_704_070_800_000_000_000), None)], + ); + assert_satisfaction( + "split evaluation errors fail closed", + &aligned, + &invalid_precision, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + let timezone_fixture = PartitioningTestFixture::new(vec![( + "timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + )])?; + let timezone_trunc = date_trunc_of(timezone_fixture.col(0), "hour", Some("UTC")); + let timezone_range = timezone_fixture.range_partitioning( + [0], + vec![timestamp_split( + Some(1_704_070_800_000_000_000), + Some("UTC"), + )], + ); + assert_satisfaction( + "timezone-aware date_trunc is not audited", + &timezone_range, + &Distribution::KeyPartitioned(vec![timezone_trunc]), + &timezone_fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + Ok(()) + } + + #[test] + fn descending_range_uses_the_open_side_successor() -> Result<()> { + let fixture = PartitioningTestFixture::new(vec![( + "timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, None), + )])?; + let hour = 1_704_070_800_000_000_000; + let ordering: LexOrdering = + [fixture.range_sort_expr(0, SortOptions::new(true, false))].into(); + let required = Distribution::KeyPartitioned(vec![date_trunc_of( + fixture.col(0), + "hour", + None, + )]); + + assert_satisfaction( + "the successor of an hour boundary is in the same bucket", + &fixture.range_partitioning_with_ordering( + ordering.clone(), + vec![timestamp_split(Some(hour), None)], + ), + &required, + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + assert_satisfaction( + "the successor of the final nanosecond starts a new bucket", + &fixture.range_partitioning_with_ordering( + ordering, + vec![timestamp_split(Some(hour + 3_600_000_000_000 - 1), None)], + ), + &required, + &fixture.eq_properties, + PartitioningSatisfaction::KeyLocal, + PartitioningSatisfaction::KeyLocal, + ); + Ok(()) + } + #[test] #[expect( deprecated, diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 776e3c3bf14bd..9324fb570fa3a 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -732,12 +732,11 @@ fn add_roundrobin_on_top( } } -// Partial aggregates require unspecified input distribution, but their output -// may already satisfy the final aggregate's key distribution because partial -// aggregation preserves/projects input partitioning. Keep that reusable output -// partitioning intact when preserve_file_partitions would otherwise insert -// RoundRobin below the partial aggregate. -fn partial_aggregate_output_satisfies_final_partitioning( +// Partial aggregates require unspecified input distribution, but their input +// may already keep every group partition-local. Preserve that layout when +// preserve_file_partitions would otherwise insert RoundRobin below the partial +// aggregate, and reuse it for the final aggregate. +fn partial_aggregate_has_partition_local_groups( plan: &Arc, allow_subset_satisfy_partitioning: bool, ) -> bool { @@ -751,15 +750,54 @@ fn partial_aggregate_output_satisfies_final_partitioning( return false; } - let key_distribution = Distribution::KeyPartitioned(aggregate.output_group_expr()); - - plan.output_partitioning() + let output_distribution = Distribution::KeyPartitioned(aggregate.output_group_expr()); + if plan + .output_partitioning() .satisfaction( - &key_distribution, + &output_distribution, plan.equivalence_properties(), allow_subset_satisfy_partitioning, ) .is_satisfied() + { + return true; + } + + let input_distribution = + Distribution::KeyPartitioned(aggregate.group_expr().input_exprs()); + aggregate + .input() + .output_partitioning() + .satisfaction( + &input_distribution, + aggregate.input().equivalence_properties(), + allow_subset_satisfy_partitioning, + ) + .is_key_local() +} + +fn final_aggregate_accepts_partition_local_groups( + final_aggregate: &Arc, + child: &Arc, + allow_subset_satisfy_partitioning: bool, +) -> bool { + let Some(final_aggregate) = final_aggregate.downcast_ref::() else { + return false; + }; + let Some(partial_aggregate) = child.downcast_ref::() else { + return false; + }; + + final_aggregate.mode() == &AggregateMode::FinalPartitioned + && !final_aggregate.group_expr().has_grouping_set() + && physical_exprs_equal( + &partial_aggregate.output_group_expr(), + &final_aggregate.group_expr().input_exprs(), + ) + && partial_aggregate_has_partition_local_groups( + child, + allow_subset_satisfy_partitioning, + ) } /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator @@ -1527,7 +1565,7 @@ pub fn ensure_distribution_with_stats( let preserve_partial_aggregate_partitioning = preserve_file_partition_threshold_met - && partial_aggregate_output_satisfies_final_partitioning( + && partial_aggregate_has_partition_local_groups( &plan, allow_subset_satisfy_partitioning, ); @@ -1570,7 +1608,12 @@ pub fn ensure_distribution_with_stats( ChildSatisfactionOptions::new() .with_allow_subset(allow_subset_satisfy_partitioning), )? - .is_satisfied(); + .is_satisfied() + || final_aggregate_accepts_partition_local_groups( + &plan, + &child.plan, + allow_subset_satisfy_partitioning, + ); let preserve_satisfying_file_partitioning = preserve_file_partition_threshold_met && !requires_grouping_id diff --git a/datafusion/physical-optimizer/src/sanity_checker.rs b/datafusion/physical-optimizer/src/sanity_checker.rs index 713213b70612d..dfb6b46fd6429 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/datafusion/physical-optimizer/src/sanity_checker.rs @@ -30,6 +30,7 @@ use datafusion_common::config::{ConfigOptions, OptimizerOptions}; use datafusion_common::plan_err; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_physical_expr::intervals::utils::{check_support, is_datatype_supported}; +use datafusion_physical_plan::aggregates::{AggregateExec, AggregateMode}; use datafusion_physical_plan::execution_plan::{ Boundedness, EmissionType, InvariantLevel, }; @@ -167,14 +168,19 @@ pub fn check_plan_sanity( } } - if !input_distributions - .child_satisfaction( - idx, - child.as_ref(), - ChildSatisfactionOptions::new().with_allow_subset(true), - )? - .is_satisfied() - { + let distribution_satisfaction = input_distributions.child_satisfaction( + idx, + child.as_ref(), + ChildSatisfactionOptions::new().with_allow_subset(true), + )?; + let accepts_partition_local_groups = plan + .downcast_ref::() + .is_some_and(|aggregate| { + aggregate.mode() == &AggregateMode::SinglePartitioned + && !aggregate.group_expr().has_grouping_set() + && distribution_satisfaction.is_key_local() + }); + if !distribution_satisfaction.is_satisfied() && !accepts_partition_local_groups { let plan_str = get_plan_string(plan); return plan_err!( "Plan: {:?} does not satisfy distribution requirements: {}. Child-{} output partitioning: {}", diff --git a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt index 18123a492dbd6..73550943d8381 100644 --- a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt +++ b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt @@ -146,6 +146,64 @@ k2 2024-01-01T00:30:00 7 k2 2024-01-01T01:30:00 30 k2 2024-01-01T01:45:00 5 +########## +# TEST 4: date_trunc('hour') is audited and the split is aligned, so each +# group is partition-local and aggregation needs no hash shuffle. +########## + +query TT +EXPLAIN SELECT key, date_trunc('hour', timestamp) AS time_bin, sum(value) +FROM range_sorted_time_bin +WHERE col4 = 'a' +GROUP BY key, time_bin; +---- +physical_plan +01)ProjectionExec: expr=[key@0 as key, date_trunc(Utf8("hour"),range_sorted_time_bin.timestamp)@1 as time_bin, sum(range_sorted_time_bin.value)@2 as sum(range_sorted_time_bin.value)] +02)--AggregateExec: mode=SinglePartitioned, gby=[key@0 as key, date_trunc(hour, timestamp@1) as date_trunc(Utf8("hour"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +03)----FilterExec: col4@1 = a, projection=[key@0, timestamp@2, value@3] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-1.parquet]]}, projection=[key, col4, timestamp, value], output_ordering=[key@0 ASC, timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], [(1704070800000000000)], 2), file_type=parquet, predicate=col4@4 = a, pruning_predicate=col4_null_count@2 != row_count@3 AND col4_min@0 <= a AND a <= col4_max@1, required_guarantees=[col4 in (a)] + +query TPI +SELECT key, date_trunc('hour', timestamp) AS time_bin, sum(value) +FROM range_sorted_time_bin +WHERE col4 = 'a' +GROUP BY key, time_bin +ORDER BY key, time_bin; +---- +k1 2024-01-01T00:00:00 3 +k1 2024-01-01T01:00:00 30 +k2 2024-01-01T00:00:00 7 +k2 2024-01-01T01:00:00 35 + +########## +# TEST 5: date_trunc('day') straddles the hour split, so the hash shuffle +# remains. +########## + +query TT +EXPLAIN SELECT key, date_trunc('day', timestamp) AS time_bin, sum(value) +FROM range_sorted_time_bin +WHERE col4 = 'a' +GROUP BY key, time_bin; +---- +physical_plan +01)ProjectionExec: expr=[key@0 as key, date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1 as time_bin, sum(range_sorted_time_bin.value)@2 as sum(range_sorted_time_bin.value)] +02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1 as date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([key@0, date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1], 2), input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1 ASC +04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_trunc(day, timestamp@1) as date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)], aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted +05)--------FilterExec: col4@1 = a, projection=[key@0, timestamp@2, value@3] +06)----------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-1.parquet]]}, projection=[key, col4, timestamp, value], output_ordering=[key@0 ASC, timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], [(1704070800000000000)], 2), file_type=parquet, predicate=col4@4 = a, pruning_predicate=col4_null_count@2 != row_count@3 AND col4_min@0 <= a AND a <= col4_max@1, required_guarantees=[col4 in (a)] + +query TPI +SELECT key, date_trunc('day', timestamp) AS time_bin, sum(value) +FROM range_sorted_time_bin +WHERE col4 = 'a' +GROUP BY key, time_bin +ORDER BY key, time_bin; +---- +k1 2024-01-01T00:00:00 33 +k2 2024-01-01T00:00:00 42 + ########## # CLEANUP ########## From dac39453a462cdcfaa434eb2aa147b832ade965f Mon Sep 17 00:00:00 2001 From: "ben.bellick" Date: Thu, 17 Sep 2026 13:01:13 +0000 Subject: [PATCH 3/4] fix(optimizer): unify aggregate range locality enforcement --- datafusion/core/Cargo.toml | 4 + datafusion/core/benches/range_date_trunc.rs | 152 +++++++++++ .../enforce_distribution.rs | 104 ++++++++ datafusion/physical-expr/src/partitioning.rs | 241 +++++++++++++++--- .../enforce_distribution.rs | 50 +--- .../physical-optimizer/src/sanity_checker.rs | 6 +- .../physical-plan/src/aggregates/mod.rs | 47 ++++ .../src/test_context/range_partitioning.rs | 6 +- .../test_files/range_sorted_time_bin_agg.slt | 23 +- 9 files changed, 532 insertions(+), 101 deletions(-) create mode 100644 datafusion/core/benches/range_date_trunc.rs diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 222c0ec688b78..7fcf5f48e6981 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -217,6 +217,10 @@ required-features = ["docs_generation"] harness = false name = "aggregate_query_sql" +[[bench]] +harness = false +name = "range_date_trunc" + [[bench]] harness = false name = "csv_load" diff --git a/datafusion/core/benches/range_date_trunc.rs b/datafusion/core/benches/range_date_trunc.rs new file mode 100644 index 0000000000000..4e7c3bc2f897d --- /dev/null +++ b/datafusion/core/benches/range_date_trunc.rs @@ -0,0 +1,152 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Group hourly range partitions without moving groups across partitions. +//! The same benchmark can run on main: only the optimized plan differs. + +use std::hint::black_box; +use std::sync::Arc; + +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion::arrow::array::{ArrayRef, Int64Array, TimestampNanosecondArray}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::catalog::streaming::StreamingTable; +use datafusion::common::{Result, ScalarValue, SplitPoint}; +use datafusion::logical_expr::col; +use datafusion::physical_expr::{ + Partitioning, PhysicalSortExpr, RangePartitioning, expressions::col as physical_col, +}; +use datafusion::physical_plan::{ + collect, execution_plan::reset_plan_states, streaming::PartitionStream, + test::TestPartitionStream, +}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use tokio::runtime::Runtime; + +const PARTITIONS: usize = 8; +const REPETITIONS: i64 = 8; + +fn context(keys_per_partition: i64) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), false), + Field::new("value", DataType::Int64, false), + ])); + let start = 1_704_067_200_000_000_000_i64; + let hour = 3_600_000_000_000_i64; + let mut streams: Vec> = vec![]; + for partition in 0..PARTITIONS { + let mut batches = vec![]; + let rows = keys_per_partition * REPETITIONS; + for offset in (0..rows).step_by(8192) { + let end = (offset + 8192).min(rows); + let columns: Vec = vec![ + Arc::new(Int64Array::from_iter_values( + (offset..end).map(|i| i / REPETITIONS), + )), + Arc::new(TimestampNanosecondArray::from_iter_values( + (offset..end).map(|i| { + start + + partition as i64 * hour + + (i % REPETITIONS) * 1_000_000_000 + }), + )), + Arc::new(Int64Array::from_iter_values((offset..end).map(|_| 1))), + ]; + batches.push(RecordBatch::try_new(Arc::clone(&schema), columns)?); + } + streams.push(Arc::new(TestPartitionStream::new_with_batches(batches))); + } + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new_default(physical_col("ts", &schema)?)].into(), + (1..PARTITIONS) + .map(|partition| { + SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(start + partition as i64 * hour), + None, + )]) + }) + .collect(), + )?); + let mut config = SessionConfig::new().with_target_partitions(PARTITIONS); + config + .options_mut() + .optimizer + .enable_round_robin_repartition = false; + config.options_mut().optimizer.subset_repartition_threshold = PARTITIONS; + let context = SessionContext::new_with_config(config); + context.register_table( + "t", + Arc::new( + StreamingTable::try_new(schema, streams)? + .with_sort_order(vec![ + col("key").sort(true, true), + col("ts").sort(true, true), + ]) + .with_output_partitioning(partitioning), + ), + )?; + Ok(context) +} + +fn criterion_benchmark(criterion: &mut Criterion) { + let runtime = Runtime::new().unwrap(); + for keys_per_partition in [1024, 65_536] { + let context = context(keys_per_partition).unwrap(); + let plan = + runtime.block_on(async { + context + .sql("SELECT key, date_trunc('hour', ts) AS bucket, sum(value) AS total \ + FROM t GROUP BY key, bucket") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap() + }); + let results = runtime + .block_on(collect(Arc::clone(&plan), context.task_ctx())) + .unwrap(); + assert_eq!( + results.iter().map(RecordBatch::num_rows).sum::(), + PARTITIONS * keys_per_partition as usize + ); + for batch in results { + let sums = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(sums.values().iter().all(|sum| *sum == REPETITIONS)); + } + criterion.bench_function( + &format!("range_date_trunc/{keys_per_partition}_keys_per_partition"), + |bencher| { + bencher.iter(|| { + let plan = reset_plan_states(Arc::clone(&plan)).unwrap(); + black_box( + runtime.block_on(collect(plan, context.task_ctx())).unwrap(), + ) + }); + }, + ); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 7c9d4984a6a68..8d342ca827199 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -402,6 +402,110 @@ fn parquet_exec_with_output_partitioning( DataSourceExec::from_data_source(config) } +#[test] +fn range_date_trunc_locality_survives_aggregate_combination() -> Result<()> { + use arrow_schema::TimeUnit; + use datafusion_physical_expr::ScalarFunctionExpr; + use datafusion_physical_optimizer::combine_partial_final_agg::CombinePartialFinalAggregate; + + let timestamp_type = DataType::Timestamp(TimeUnit::Nanosecond, None); + let schema = Arc::new(Schema::new(vec![Field::new( + "ts", + timestamp_type.clone(), + true, + )])); + let timestamp = col("ts", &schema)?; + let range = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new_default(Arc::clone(×tamp))].into(), + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(1_704_070_800_000_000_000), + None, + )])], + )?); + let scan = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///")?, + Arc::new(ParquetSource::new(Arc::clone(&schema))), + ) + .with_file_groups(vec![ + FileGroup::new(vec![PartitionedFile::new("p0", 100)]), + FileGroup::new(vec![PartitionedFile::new("p1", 100)]), + ]) + .with_output_partitioning(Some(range)) + .build(); + let source: Arc = DataSourceExec::from_data_source(scan); + let bucket: Arc = Arc::new(ScalarFunctionExpr::new( + "date_trunc", + datafusion_functions::datetime::date_trunc(), + vec![lit("hour"), timestamp], + Field::new("bucket", timestamp_type, true).into(), + Arc::new(ConfigOptions::default()), + )); + let group_by = + PhysicalGroupBy::new_single(vec![(Arc::clone(&bucket), "bucket".into())]); + let partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + Arc::clone(&source), + Arc::clone(&schema), + )?; + let final_plan: Arc = Arc::new(AggregateExec::try_new( + AggregateMode::FinalPartitioned, + group_by.as_final(), + vec![], + vec![], + Arc::new(partial), + schema, + )?); + let mut config = ConfigOptions::default(); + config.execution.target_partitions = 2; + config.optimizer.enable_round_robin_repartition = false; + config.optimizer.repartition_file_scans = false; + let enforced = EnsureRequirements::new().optimize(final_plan, &config)?; + let plan_text = displayable(enforced.as_ref()).indent(true).to_string(); + assert!(plan_text.contains("mode=FinalPartitioned"), "{plan_text}"); + assert!(plan_text.contains("mode=Partial"), "{plan_text}"); + assert!(!plan_text.contains("RepartitionExec"), "{plan_text}"); + SanityCheckPlan::new().optimize(Arc::clone(&enforced), &config)?; + + let combined = CombinePartialFinalAggregate::new().optimize(enforced, &config)?; + let combined_text = displayable(combined.as_ref()).indent(true).to_string(); + assert!( + combined_text.contains("mode=SinglePartitioned"), + "{combined_text}" + ); + let reenforced = EnsureRequirements::new().optimize(combined, &config)?; + assert_eq!( + displayable(reenforced.as_ref()).indent(true).to_string(), + combined_text + ); + SanityCheckPlan::new().optimize(reenforced, &config)?; + + // Group locality does not establish a join-compatible partition layout. + let join = crate::physical_optimizer::test_utils::hash_join_exec( + Arc::clone(&source), + source, + vec![(Arc::clone(&bucket), bucket)], + None, + &JoinType::Inner, + )?; + assert!( + SanityCheckPlan::new() + .optimize(Arc::clone(&join), &config) + .is_err() + ); + let enforced_join = EnsureRequirements::new().optimize(join, &config)?; + assert!( + displayable(enforced_join.as_ref()) + .indent(true) + .to_string() + .contains("partitioning=Hash") + ); + SanityCheckPlan::new().optimize(enforced_join, &config)?; + Ok(()) +} + fn range_partitioning( column: &str, split_values: impl IntoIterator, diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 6ddcd4df82a47..91326d1cf062d 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -421,42 +421,39 @@ fn evaluate_with_range_value( } } -fn projected_split_points( +/// For a nondecreasing, null-preserving transform, adjacent partition images +/// are disjoint when the values immediately across each split remain distinct. +/// The open side is below an ASC split and above a DESC split. +fn range_transform_keeps_groups_local( function: &Arc, range_key: &Arc, range: &RangePartitioning, key_index: usize, schema: &Schema, -) -> Option> { +) -> Option<()> { let range_type = range_key.data_type(schema).ok()?; - range - .split_points() - .iter() - .map(|split_point| { - let split_value = split_point.values().get(key_index)?; - if split_value.data_type() != range_type { - return None; - } - let adjacent = if range.ordering()[key_index].options.descending { - checked_successor(split_value) - } else { - checked_predecessor(split_value) - }?; - let at_split = evaluate_with_range_value(function, range_key, split_value)?; - let across_split = evaluate_with_range_value(function, range_key, &adjacent)?; - if at_split.is_null() - || across_split.is_null() - || at_split.data_type() != across_split.data_type() - || at_split == across_split - { - return None; - } + range.split_points().iter().try_for_each(|split_point| { + let split_value = split_point.values().get(key_index)?; + if split_value.data_type() != range_type { + return None; + } + let adjacent = if range.ordering()[key_index].options.descending { + checked_successor(split_value) + } else { + checked_predecessor(split_value) + }?; + let at_split = evaluate_with_range_value(function, range_key, split_value)?; + let across_split = evaluate_with_range_value(function, range_key, &adjacent)?; + if at_split.is_null() + || across_split.is_null() + || at_split.data_type() != across_split.data_type() + || at_split == across_split + { + return None; + } - let mut values = split_point.values().to_vec(); - values[key_index] = at_split; - Some(SplitPoint::new(values)) - }) - .collect() + Some(()) + }) } fn is_audited_range_transform( @@ -470,31 +467,29 @@ fn is_audited_range_transform( return false; }; let schema = eq_properties.schema(); - if !function - .supports_range_partitioning_analysis(schema) - .unwrap_or(false) + if datafusion_physical_expr_common::physical_expr::is_volatile(candidate) + || !function + .supports_range_partitioning_analysis(schema) + .unwrap_or(false) || !eq_properties.is_same_direction_ordered_transform(candidate, range_key) { return false; } let mut range_argument_count = 0; - let dummy_batch = create_dummy_batch().ok(); for argument in function.args() { if equivalent_expr(argument, range_key, eq_properties) { range_argument_count += 1; - } else if dummy_batch - .as_ref() - .and_then(|batch| argument.evaluate(batch).ok()) - .is_none_or(|value| !matches!(value, ColumnarValue::Scalar(_))) - { + } else if !argument.is::() { return false; } } range_argument_count == 1 - && projected_split_points(candidate, range_key, range, key_index, schema) - .is_some() + && range_transform_keeps_groups_local( + candidate, range_key, range, key_index, schema, + ) + .is_some() } fn range_transform_satisfaction( @@ -1196,6 +1191,172 @@ mod tests { )]) } + fn assert_date_trunc_partition_images() + -> Result<()> { + use arrow::array::{PrimitiveArray, RecordBatch}; + use datafusion_common::HashSet; + + let ticks_per_second = match T::UNIT { + TimeUnit::Second => 1, + TimeUnit::Millisecond => 1_000, + TimeUnit::Microsecond => 1_000_000, + TimeUnit::Nanosecond => 1_000_000_000, + }; + let hour = 3600 * ticks_per_second; + let data_type = DataType::Timestamp(T::UNIT, None); + let schema = + Arc::new(Schema::new(vec![Field::new("ts", data_type.clone(), true)])); + let timestamp: Arc = Arc::new(Column::new("ts", 0)); + let properties = EquivalenceProperties::new(Arc::clone(&schema)); + + let mut cases = vec![ + ("hour", false, vec![-hour, 0, hour], true), + ("hour", false, vec![-hour + 1, 1, hour + 1], false), + ("day", false, vec![-hour, 0, hour], false), + ("hour", true, vec![-hour, 0, hour], false), + ("hour", true, vec![-hour - 1, -1, hour - 1], true), + ]; + for (precision, divisor) in [("millisecond", 1_000), ("microsecond", 1_000_000)] { + let width = (ticks_per_second / divisor).max(1); + cases.extend([ + (precision, false, vec![-width, 0, width], true), + // In particular, microsecond input with millisecond precision + // must reject split -999: both adjacent array values map to -1000. + (precision, false, vec![-width + 1], width == 1), + (precision, true, vec![-width, 0, width], width == 1), + (precision, true, vec![-width - 1, -1, width - 1], true), + ]); + } + for (precision, descending, source_splits, expect_disjoint) in cases { + for nulls_first in [false, true] { + let options = SortOptions::new(descending, nulls_first); + let mut splits = source_splits.clone(); + if descending { + splits.reverse(); + } + let range = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new(Arc::clone(×tamp), options)].into(), + splits + .iter() + .map(|value| { + SplitPoint::new(vec![ScalarValue::new_timestamp::( + Some(*value), + None, + )]) + }) + .collect(), + )?); + let function: Arc = Arc::new(ScalarFunctionExpr::new( + "date_trunc", + datafusion_functions::datetime::date_trunc(), + vec![ + Arc::new(Literal::new(ScalarValue::from(precision))), + Arc::clone(×tamp), + ], + Field::new("bucket", data_type.clone(), true).into(), + Arc::new(ConfigOptions::default()), + )); + let required = Distribution::KeyPartitioned(vec![Arc::clone(&function)]); + assert_eq!( + range + .satisfaction(&required, &properties, false) + .is_key_local(), + expect_disjoint + ); + + // Assign rows by source-domain comparisons, independently of + // the optimizer's adjacency check, then evaluate whole arrays. + let mut partitions = vec![vec![]; splits.len() + 1]; + let rows = splits + .iter() + .flat_map(|split| [Some(split - 1), Some(*split), Some(split + 1)]) + .chain([None, None]); + for value in rows { + let index = match value { + None if nulls_first => 0, + None => splits.len(), + Some(value) => splits + .iter() + .filter(|split| { + if descending { + value <= **split + } else { + value >= **split + } + }) + .count(), + }; + partitions[index].push(value); + } + let images = partitions + .into_iter() + .map(|values| { + let input = Arc::new(PrimitiveArray::::from_iter(values)); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![input])?; + let output = + function.evaluate(&batch)?.into_array(batch.num_rows())?; + (0..output.len()) + .map(|i| ScalarValue::try_from_array(&output, i)) + .collect::>>() + }) + .collect::>>()?; + let disjoint = images.iter().enumerate().all(|(i, left)| { + images + .iter() + .skip(i + 1) + .all(|right| left.is_disjoint(right)) + }); + assert_eq!( + disjoint, + expect_disjoint, + "{precision}, {:?}, {options:?}, splits={splits:?}", + T::UNIT + ); + } + } + Ok(()) + } + + #[test] + fn date_trunc_partition_images_are_disjoint_only_for_safe_splits() -> Result<()> { + use arrow::array::types::{ + TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, + TimestampSecondType, + }; + assert_date_trunc_partition_images::()?; + assert_date_trunc_partition_images::()?; + assert_date_trunc_partition_images::()?; + assert_date_trunc_partition_images::() + } + + #[test] + fn range_transform_requires_literal_parameters() -> Result<()> { + use crate::expressions::CastExpr; + let fixture = PartitioningTestFixture::new(vec![( + "ts", + DataType::Timestamp(TimeUnit::Nanosecond, None), + )])?; + let function = date_trunc_of(fixture.col(0), "hour", None); + let scalar_but_not_literal: Arc = Arc::new(CastExpr::new( + Arc::new(Literal::new(ScalarValue::from("hour"))), + DataType::Utf8View, + None, + )); + let function = + function.with_new_children(vec![scalar_but_not_literal, fixture.col(0)])?; + let range = fixture.range_partitioning([0], vec![timestamp_split(Some(0), None)]); + assert_eq!( + range.satisfaction( + &Distribution::KeyPartitioned(vec![function]), + &fixture.eq_properties, + false, + ), + PartitioningSatisfaction::NotSatisfied + ); + Ok(()) + } + #[test] fn range_transform_satisfaction_is_key_local_only() -> Result<()> { let fixture = PartitioningTestFixture::new(vec![ diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 9324fb570fa3a..6a124d2a4c7e5 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -763,41 +763,10 @@ fn partial_aggregate_has_partition_local_groups( return true; } - let input_distribution = - Distribution::KeyPartitioned(aggregate.group_expr().input_exprs()); - aggregate - .input() - .output_partitioning() - .satisfaction( - &input_distribution, - aggregate.input().equivalence_properties(), - allow_subset_satisfy_partitioning, - ) - .is_key_local() -} - -fn final_aggregate_accepts_partition_local_groups( - final_aggregate: &Arc, - child: &Arc, - allow_subset_satisfy_partitioning: bool, -) -> bool { - let Some(final_aggregate) = final_aggregate.downcast_ref::() else { - return false; - }; - let Some(partial_aggregate) = child.downcast_ref::() else { - return false; - }; - - final_aggregate.mode() == &AggregateMode::FinalPartitioned - && !final_aggregate.group_expr().has_grouping_set() - && physical_exprs_equal( - &partial_aggregate.output_group_expr(), - &final_aggregate.group_expr().input_exprs(), - ) - && partial_aggregate_has_partition_local_groups( - child, - allow_subset_satisfy_partitioning, - ) + aggregate.groups_are_partition_local( + aggregate.input().as_ref(), + allow_subset_satisfy_partitioning, + ) } /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator @@ -1609,10 +1578,13 @@ pub fn ensure_distribution_with_stats( .with_allow_subset(allow_subset_satisfy_partitioning), )? .is_satisfied() - || final_aggregate_accepts_partition_local_groups( - &plan, - &child.plan, - allow_subset_satisfy_partitioning, + || plan.downcast_ref::().is_some_and( + |aggregate| { + aggregate.groups_are_partition_local( + child.plan.as_ref(), + allow_subset_satisfy_partitioning, + ) + }, ); let preserve_satisfying_file_partitioning = preserve_file_partition_threshold_met diff --git a/datafusion/physical-optimizer/src/sanity_checker.rs b/datafusion/physical-optimizer/src/sanity_checker.rs index dfb6b46fd6429..0d2138014f7ee 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/datafusion/physical-optimizer/src/sanity_checker.rs @@ -30,7 +30,7 @@ use datafusion_common::config::{ConfigOptions, OptimizerOptions}; use datafusion_common::plan_err; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_physical_expr::intervals::utils::{check_support, is_datatype_supported}; -use datafusion_physical_plan::aggregates::{AggregateExec, AggregateMode}; +use datafusion_physical_plan::aggregates::AggregateExec; use datafusion_physical_plan::execution_plan::{ Boundedness, EmissionType, InvariantLevel, }; @@ -176,9 +176,7 @@ pub fn check_plan_sanity( let accepts_partition_local_groups = plan .downcast_ref::() .is_some_and(|aggregate| { - aggregate.mode() == &AggregateMode::SinglePartitioned - && !aggregate.group_expr().has_grouping_set() - && distribution_satisfaction.is_key_local() + aggregate.groups_are_partition_local(child.as_ref(), true) }); if !distribution_satisfaction.is_satisfied() && !accepts_partition_local_groups { let plan_str = get_plan_string(plan); diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index ffa93a387d87a..cdfa9bbd0eff1 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1212,6 +1212,53 @@ impl AggregateExec { &self.input } + /// Whether a candidate input keeps each group in one partition. + /// + /// Unlike co-partitioning, aggregation only needs equality groups to be + /// local. A final stage can inherit that guarantee from a partial stage + /// even when projecting the grouping expressions loses range metadata. + /// This check does not require the two stages to be combined. + pub fn groups_are_partition_local( + &self, + input: &dyn ExecutionPlan, + allow_subset: bool, + ) -> bool { + if !matches!( + self.mode, + AggregateMode::Partial + | AggregateMode::SinglePartitioned + | AggregateMode::FinalPartitioned + ) || self.group_by.is_empty() + || self.group_by.has_grouping_set() + { + return false; + } + + let required = Distribution::KeyPartitioned(self.group_by.input_exprs()); + if input + .output_partitioning() + .satisfaction(&required, input.equivalence_properties(), allow_subset) + .is_key_local() + { + return true; + } + + self.mode == AggregateMode::FinalPartitioned + && input + .downcast_ref::() + .is_some_and(|partial| { + partial.mode == AggregateMode::Partial + && datafusion_physical_expr::physical_exprs_equal( + &partial.output_group_expr(), + &self.group_by.input_exprs(), + ) + && partial.groups_are_partition_local( + partial.input.as_ref(), + allow_subset, + ) + }) + } + /// Get the input schema before any aggregates are applied pub fn input_schema(&self) -> SchemaRef { Arc::clone(&self.input_schema) diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index becde0f3286db..739534cde542d 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -344,9 +344,9 @@ type TimeBinRow = ( /// - partition 1: `[2024-01-01 01:00, 02:00)` /// /// Files are range-partitioned on `timestamp` and sorted on `(key, timestamp)`. -/// Because `date_bin(60 seconds, timestamp)` does not straddle the hour split, -/// grouping by `(key, time_bin)` is partition-disjoint. Today's planner still -/// inserts a hash shuffle; the test pins that plan so a follow-up can remove it. +/// The hour split permits partition-local `date_trunc('hour')` aggregation, +/// but `date_trunc('day')` straddles it. Aligned `date_bin` remains a negative +/// control because the function is not admitted to range analysis. pub(super) fn register_range_sorted_time_bin_table(ctx: &SessionContext) { let schema = Arc::new(Schema::new(vec![ Field::new("key", DataType::Utf8, false), diff --git a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt index 73550943d8381..b480db12da0ce 100644 --- a/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt +++ b/datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt @@ -29,15 +29,10 @@ # 1. Range([timestamp]) and output_ordering=[key, timestamp] # 2. Two file_groups, so the two 60-minute streams run in parallel # -# Improvement opportunity: -# date_bin(60s) is monotonic in timestamp and the hour split is aligned to bin -# boundaries, so (key, time_bin) is partition-disjoint. Aggregation could be a -# single streaming SinglePartitioned step with no hash shuffle. -# -# Today's plan still hash-repartitions: -# Partial AggregateExec (ordering_mode=Sorted) -# -> RepartitionExec Hash([key, date_bin(...)]) -# -> FinalPartitioned AggregateExec (ordering_mode=Sorted) +# date_bin is a negative control: even aligned bins retain the hash shuffle +# because date_bin can produce NULL from non-null inputs and is not admitted +# to range analysis. Timezone-less date_trunc is admitted, but only removes +# the shuffle when every split keeps the resulting groups partition-local. statement ok set datafusion.explain.physical_plan_only = true; @@ -82,11 +77,9 @@ physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion ########## # TEST 2: Filtered time-bin aggregation. # GROUP BY keys are (key, date_bin(timestamp)). Input is sorted on those keys -# (date_bin is monotonic in timestamp) and range-partitioned so bins do not -# overlap across the two 60-minute streams. -# -# Today this is still Partial + hash RepartitionExec + Final, even though -# ordering_mode=Sorted is already recognized. +# and the split aligns with the bins in this data. The function remains +# default-denied by range analysis, so Partial + hash RepartitionExec + Final +# is retained even though ordering_mode=Sorted is recognized. ########## query TT @@ -118,7 +111,7 @@ k2 2024-01-01T01:45:00 5 ########## # TEST 3: Same aggregation without the col4 filter. The scan still has two -# 60-minute file groups, and today's plan still hash-repartitions. +# 60-minute file groups, and date_bin still requires a hash repartition. ########## query TT From 35ea93952e86602c6b305da12e28ae4651f02ab9 Mon Sep 17 00:00:00 2001 From: "ben.bellick" Date: Fri, 18 Sep 2026 17:40:36 +0000 Subject: [PATCH 4/4] test(datetime): strengthen date_trunc range audit --- .../functions/src/datetime/date_trunc.rs | 141 ++++++++++++++---- 1 file changed, 109 insertions(+), 32 deletions(-) diff --git a/datafusion/functions/src/datetime/date_trunc.rs b/datafusion/functions/src/datetime/date_trunc.rs index 39ecaa06ad782..eef7c37f91df7 100644 --- a/datafusion/functions/src/datetime/date_trunc.rs +++ b/datafusion/functions/src/datetime/date_trunc.rs @@ -908,8 +908,8 @@ mod tests { use arrow::compute::DatePart; use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos; use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit}; - use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; + use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; #[test] @@ -968,19 +968,17 @@ mod tests { timestamp: ColumnarValue, timestamp_type: DataType, number_rows: usize, - ) -> ColumnarValue { - DateTruncFunc::new() - .invoke_with_args(ScalarFunctionArgs { - args: vec![ScalarValue::from(granularity).into(), timestamp], - arg_fields: vec![ - Field::new("granularity", DataType::Utf8, false).into(), - Field::new("timestamp", timestamp_type.clone(), true).into(), - ], - number_rows, - return_field: Field::new("date_trunc", timestamp_type, true).into(), - config_options: Arc::new(ConfigOptions::default()), - }) - .unwrap() + ) -> Result { + DateTruncFunc::new().invoke_with_args(ScalarFunctionArgs { + args: vec![ScalarValue::from(granularity).into(), timestamp], + arg_fields: vec![ + Field::new("granularity", DataType::Utf8, false).into(), + Field::new("timestamp", timestamp_type.clone(), true).into(), + ], + number_rows, + return_field: Field::new("date_trunc", timestamp_type, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }) } fn assert_timezone_less_range_analysis_contract() { @@ -990,22 +988,41 @@ mod tests { TimeUnit::Microsecond => 1_000, TimeUnit::Nanosecond => 1, }; - let values = [ - Some("1969-12-28T23:59:59Z"), - Some("1969-12-31T23:59:59Z"), - Some("1970-01-01T00:00:00Z"), - Some("1999-12-31T23:59:59Z"), - Some("2000-01-01T00:00:00Z"), - Some("2000-02-28T23:59:59Z"), - Some("2000-02-29T00:00:00Z"), - Some("2000-03-01T00:00:00Z"), - Some("2023-12-31T23:59:59Z"), - Some("2024-01-01T00:00:00Z"), - None, + // Probe both sides of calendar and fixed-width boundaries at the + // input's own resolution, including negative fractional timestamps. + let mut timestamps = [ + "1900-03-01T00:00:00Z", // Non-leap century. + "1969-12-29T00:00:00Z", // Monday before the epoch. + "1970-01-01T00:00:00Z", + "1970-01-01T00:00:00.000001Z", + "1970-01-01T00:00:00.001Z", + "1970-01-01T00:00:01Z", + "1970-01-01T00:01:00Z", + "1970-01-01T01:00:00Z", + "2000-01-01T00:00:00Z", + "2000-02-29T00:00:00Z", // Leap century. + "2000-03-01T00:00:00Z", + "2000-04-01T00:00:00Z", // Quarter boundary. + "2024-01-01T00:00:00Z", // Monday and year boundary. ] - .map(|value| { - value.map(|value| string_to_timestamp_nanos(value).unwrap() / scale) - }); + .into_iter() + .flat_map(|value| { + let boundary = string_to_timestamp_nanos(value).unwrap() / scale; + [boundary - 1, boundary, boundary + 1] + }) + .collect::>(); + // Stay inside the scalar nanosecond domain with room to truncate a + // full year at its lower edge; the upper edge needs no such margin. + timestamps.extend([ + i64::MIN / scale + 366 * 86_400 * (1_000_000_000 / scale), + i64::MAX / scale, + ]); + timestamps.sort_unstable(); + timestamps.dedup(); + let values = timestamps + .into_iter() + .flat_map(|value| [Some(value), None]) + .collect::>(); let timestamp_type = DataType::Timestamp(T::UNIT, None); for granularity in [ @@ -1020,13 +1037,14 @@ mod tests { "quarter", "year", ] { - let input = PrimitiveArray::::from_iter(values); + let input = PrimitiveArray::::from_iter(values.iter().copied()); let output = invoke_date_trunc( granularity, ColumnarValue::Array(Arc::new(input.clone())), timestamp_type.clone(), input.len(), - ); + ) + .unwrap(); let ColumnarValue::Array(output) = output else { panic!("expected array output") }; @@ -1057,7 +1075,8 @@ mod tests { ColumnarValue::Scalar(ScalarValue::new_timestamp::(*value, None)), timestamp_type.clone(), 1, - ); + ) + .unwrap(); let ColumnarValue::Scalar(scalar_output) = scalar_output else { panic!("expected scalar output") }; @@ -1079,6 +1098,64 @@ mod tests { assert_timezone_less_range_analysis_contract::(); } + #[test] + fn range_analysis_distinguishes_errors_from_null_results() { + let timestamp_type = DataType::Timestamp(TimeUnit::Nanosecond, None); + for (precision, value) in [ + ("microsecond", i64::MIN), + ("millisecond", i64::MIN), + ("hour", i64::MIN), + ("invalid", 0), + ] { + for input in [ + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond( + Some(value), + None, + )), + ColumnarValue::Array(Arc::new(TimestampNanosecondArray::from(vec![ + Some(value), + ]))), + ] { + assert!( + invoke_date_trunc(precision, input, timestamp_type.clone(), 1) + .is_err() + ); + } + } + + // The scalar path scales to nanoseconds, whereas the array fast path + // works in the original unit. A failed scalar boundary evaluation must + // decline optimization, not imply that array execution produces NULL. + let seconds = DataType::Timestamp(TimeUnit::Second, None); + assert!( + invoke_date_trunc( + "second", + ScalarValue::TimestampSecond(Some(i64::MAX), None).into(), + seconds.clone(), + 1, + ) + .is_err() + ); + let output = invoke_date_trunc( + "second", + ColumnarValue::Array(Arc::new(TimestampSecondArray::from(vec![ + Some(i64::MAX), + None, + ]))), + seconds, + 2, + ) + .unwrap(); + let ColumnarValue::Array(output) = output else { + panic!("expected array") + }; + assert_eq!( + ScalarValue::try_from_array(&output, 0).unwrap(), + ScalarValue::TimestampSecond(Some(i64::MAX), None) + ); + assert!(output.is_null(1)); + } + #[test] fn date_trunc_test() { let cases = vec![