From 6c4ece1ba162bd3a27e45c313d1467c4057398fd Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Mon, 13 Jul 2026 20:22:45 +0000 Subject: [PATCH] wip --- datafusion/physical-expr/src/lib.rs | 1 + datafusion/physical-expr/src/partitioning.rs | 532 +++++++++++++++++- datafusion/physical-optimizer/src/utils.rs | 48 +- .../src/distribution_requirements.rs | 54 +- .../src/test_context/range_partitioning.rs | 111 +++- .../test_files/range_partitioning.slt | 121 ++++ 6 files changed, 759 insertions(+), 108 deletions(-) diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index 80e9f88b510ed..de2c02d6b0378 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -65,6 +65,7 @@ pub use equivalence::{ pub use expressions::{DynamicFilterTracker, DynamicFilterTracking}; pub use partitioning::{ Distribution, Partitioning, PartitioningSatisfaction, RangePartitioning, + range_partitioning_satisfaction_for_key_partitioning, }; pub use physical_expr::{ add_offset_to_expr, add_offset_to_physical_sort_exprs, create_lex_ordering, diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 61492934ebd19..394e97b3f301a 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -18,17 +18,26 @@ //! [`Partitioning`] and [`Distribution`] for `ExecutionPlans` use crate::{ - EquivalenceProperties, PhysicalExpr, equivalence::ProjectionMapping, - expressions::UnKnownColumn, physical_exprs_contains, physical_exprs_equal, + EquivalenceProperties, PhysicalExpr, ScalarFunctionExpr, + equivalence::ProjectionMapping, + expressions::{Literal, UnKnownColumn}, + physical_exprs_contains, physical_exprs_equal, }; +use arrow::array::types::{IntervalDayTimeType, IntervalMonthDayNanoType}; +use arrow::datatypes::DataType; pub use datafusion_common::SplitPoint; -use datafusion_common::{Result, validate_range_split_points}; +use datafusion_common::{Result, ScalarValue, validate_range_split_points}; use datafusion_physical_expr_common::physical_expr::format_physical_expr_list; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use std::fmt; use std::fmt::Display; use std::sync::Arc; +const NANOS_PER_MICRO: i64 = 1_000; +const NANOS_PER_MILLI: i64 = 1_000_000; +const NANOS_PER_SECOND: i64 = 1_000_000_000; +const NANOS_PER_DAY: i128 = 86_400_000_000_000; + /// Output partitioning supported by [`ExecutionPlan`]s. /// /// Calling [`ExecutionPlan::execute`] produce one or more independent streams of @@ -260,23 +269,248 @@ impl RangePartitioning { .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 ordering = LexOrdering::new(sort_exprs)?; - if ordering.len() != self.ordering.len() { + .collect::>>(); + + if let Some(projected_exprs) = projected_exprs { + let sort_exprs = self + .ordering + .iter() + .zip(projected_exprs) + .map(|(sort_expr, expr)| PhysicalSortExpr::new(expr, sort_expr.options)) + .collect::>(); + let ordering = LexOrdering::new(sort_exprs)?; + if ordering.len() != self.ordering.len() { + return None; + } + + return Some(Self { + ordering, + split_points: self.split_points.clone(), + }); + } + + self.project_date_bin(mapping, input_eq_properties) + } + + fn project_date_bin( + &self, + mapping: &ProjectionMapping, + input_eq_properties: &EquivalenceProperties, + ) -> Option { + let leading_sort_expr = self.ordering.first(); + let leading_expr = &leading_sort_expr.expr; + + for (source_expr, targets) in mapping.iter() { + let Some(date_bin) = DateBinPartitionExpr::try_new(source_expr) else { + continue; + }; + if !exprs_equivalent(leading_expr, &date_bin.source_expr, input_eq_properties) + { + continue; + } + + let split_points = + date_bin_projected_split_points(self, &date_bin, input_eq_properties)?; + let (target_expr, _) = targets.first(); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::clone(target_expr), + leading_sort_expr.options, + )])?; + + return RangePartitioning::try_new(ordering, split_points).ok(); + } + + None + } +} + +struct DateBinPartitionExpr { + source_expr: Arc, + stride_nanos: i64, + origin_nanos: i64, +} + +impl DateBinPartitionExpr { + fn try_new(expr: &Arc) -> Option { + let func = expr.downcast_ref::()?; + if func.name() != "date_bin" { return None; } + let args = func.args(); + if !(2..=3).contains(&args.len()) { + return None; + } + + let stride_nanos = literal_value(&args[0]).and_then(stride_to_nanos)?; + if stride_nanos <= 0 { + return None; + } + + let origin_nanos = if args.len() == 3 { + literal_value(&args[2]).and_then(timestamp_to_nanos)? + } else { + 0 + }; + Some(Self { - ordering, - split_points: self.split_points.clone(), + source_expr: Arc::clone(&args[1]), + stride_nanos, + origin_nanos, }) } + + fn project_boundary(&self, value: &ScalarValue) -> Option { + let binned = date_bin_timestamp(value, self.stride_nanos, self.origin_nanos)?; + (binned == *value).then_some(binned) + } +} + +fn literal_value(expr: &Arc) -> Option<&ScalarValue> { + expr.downcast_ref::().map(Literal::value) +} + +fn stride_to_nanos(value: &ScalarValue) -> Option { + let nanos = match value { + ScalarValue::IntervalDayTime(Some(value)) => { + let (days, millis) = IntervalDayTimeType::to_parts(*value); + i128::from(days) * NANOS_PER_DAY + + i128::from(millis) * i128::from(NANOS_PER_MILLI) + } + ScalarValue::IntervalMonthDayNano(Some(value)) => { + let (months, days, nanos) = IntervalMonthDayNanoType::to_parts(*value); + if months != 0 { + return None; + } + i128::from(days) * NANOS_PER_DAY + i128::from(nanos) + } + _ => return None, + }; + + i64::try_from(nanos).ok() +} + +fn timestamp_to_nanos(value: &ScalarValue) -> Option { + match value { + ScalarValue::TimestampSecond(Some(value), _) => { + value.checked_mul(NANOS_PER_SECOND) + } + ScalarValue::TimestampMillisecond(Some(value), _) => { + value.checked_mul(NANOS_PER_MILLI) + } + ScalarValue::TimestampMicrosecond(Some(value), _) => { + value.checked_mul(NANOS_PER_MICRO) + } + ScalarValue::TimestampNanosecond(Some(value), _) => Some(*value), + _ => None, + } +} + +fn date_bin_timestamp( + value: &ScalarValue, + stride_nanos: i64, + origin_nanos: i64, +) -> Option { + let source_nanos = timestamp_to_nanos(value)?; + let binned_nanos = date_bin_nanos(stride_nanos, source_nanos, origin_nanos)?; + + match value { + ScalarValue::TimestampSecond(_, timezone) => Some(ScalarValue::TimestampSecond( + Some(binned_nanos / NANOS_PER_SECOND), + timezone.clone(), + )), + ScalarValue::TimestampMillisecond(_, timezone) => { + Some(ScalarValue::TimestampMillisecond( + Some(binned_nanos / NANOS_PER_MILLI), + timezone.clone(), + )) + } + ScalarValue::TimestampMicrosecond(_, timezone) => { + Some(ScalarValue::TimestampMicrosecond( + Some(binned_nanos / NANOS_PER_MICRO), + timezone.clone(), + )) + } + ScalarValue::TimestampNanosecond(_, timezone) => Some( + ScalarValue::TimestampNanosecond(Some(binned_nanos), timezone.clone()), + ), + _ => None, + } +} + +fn date_bin_nanos(stride: i64, source: i64, origin: i64) -> Option { + if stride <= 0 { + return None; + } + + let time_diff = source.checked_sub(origin)?; + let remainder = time_diff.checked_rem(stride)?; + let time_delta = time_diff.checked_sub(remainder)?; + let time_delta = if time_diff < 0 && stride > 1 && time_delta != time_diff { + time_delta.checked_sub(stride)? + } else { + time_delta + }; + + origin.checked_add(time_delta) +} + +fn date_bin_projected_split_points( + range: &RangePartitioning, + date_bin: &DateBinPartitionExpr, + eq_properties: &EquivalenceProperties, +) -> Option> { + let schema = eq_properties.schema(); + let ordering = range.ordering(); + let leading_sort_expr = ordering.first(); + if leading_sort_expr.options.descending + || leading_sort_expr + .expr + .nullable(schema.as_ref()) + .ok() + .unwrap_or(true) + { + return None; + } + + let leading_type = leading_sort_expr.expr.data_type(schema.as_ref()).ok()?; + if !matches!(leading_type, DataType::Timestamp(_, _)) { + return None; + } + + let trailing_min_values = ordering + .iter() + .skip(1) + .map(|sort_expr| { + if sort_expr.options.descending + || sort_expr + .expr + .nullable(schema.as_ref()) + .ok() + .unwrap_or(true) + { + return None; + } + let data_type = sort_expr.expr.data_type(schema.as_ref()).ok()?; + ScalarValue::min(&data_type) + }) + .collect::>>()?; + + range + .split_points() + .iter() + .map(|split_point| { + let values = split_point.values(); + let first = values.first()?; + let projected = date_bin.project_boundary(first)?; + for (value, min_value) in values.iter().skip(1).zip(&trailing_min_values) { + if value != min_value { + return None; + } + } + Some(SplitPoint::new(vec![projected])) + }) + .collect() } impl Display for RangePartitioning { @@ -320,6 +554,18 @@ fn equivalent_exprs( physical_exprs_equal(&normalized_left, &normalized_right) } +fn exprs_equivalent( + left: &Arc, + right: &Arc, + eq_properties: &EquivalenceProperties, +) -> bool { + equivalent_exprs( + std::slice::from_ref(left), + std::slice::from_ref(right), + eq_properties, + ) +} + fn normalize_exprs( exprs: &[Arc], eq_properties: &EquivalenceProperties, @@ -331,6 +577,84 @@ fn normalize_exprs( .collect() } +/// Returns how [`Partitioning::Range`] satisfies a key partitioning +/// requirement. +/// +/// This is intentionally separate from [`Partitioning::satisfaction`] while +/// range reuse is rolled out operator by operator. In addition to exact and +/// subset expression checks, this recognizes `date_bin` keys when all range +/// split points are aligned to the bin boundaries. +pub fn range_partitioning_satisfaction_for_key_partitioning( + partitioning: &Partitioning, + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, + allow_subset: bool, +) -> PartitioningSatisfaction { + let Partitioning::Range(range) = partitioning else { + return PartitioningSatisfaction::NotSatisfied; + }; + + let partition_exprs = range + .ordering() + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::>(); + + if partition_exprs.is_empty() || required_exprs.is_empty() { + return PartitioningSatisfaction::NotSatisfied; + } + + let normalized_partition_exprs = normalize_exprs(&partition_exprs, eq_properties); + let normalized_required_exprs = normalize_exprs(required_exprs, eq_properties); + + if physical_exprs_equal(&normalized_required_exprs, &normalized_partition_exprs) { + return PartitioningSatisfaction::Exact; + } + + if allow_subset + && normalized_partition_exprs.len() < normalized_required_exprs.len() + && normalized_partition_exprs.iter().all(|partition_expr| { + normalized_required_exprs + .iter() + .any(|required_expr| partition_expr.eq(required_expr)) + }) + { + return PartitioningSatisfaction::Subset; + } + + range_date_bin_satisfaction(range, required_exprs, eq_properties, allow_subset) + .unwrap_or(PartitioningSatisfaction::NotSatisfied) +} + +fn range_date_bin_satisfaction( + range: &RangePartitioning, + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, + allow_subset: bool, +) -> Option { + let leading_expr = &range.ordering().first().expr; + + for required_expr in required_exprs { + let Some(date_bin) = DateBinPartitionExpr::try_new(required_expr) else { + continue; + }; + if !exprs_equivalent(leading_expr, &date_bin.source_expr, eq_properties) { + continue; + } + date_bin_projected_split_points(range, &date_bin, eq_properties)?; + + return if required_exprs.len() == 1 { + Some(PartitioningSatisfaction::Exact) + } else if allow_subset { + Some(PartitioningSatisfaction::Subset) + } else { + Some(PartitioningSatisfaction::NotSatisfied) + }; + } + + None +} + /// Represents how a [`Partitioning`] satisfies a [`Distribution`] requirement. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PartitioningSatisfaction { @@ -572,9 +896,12 @@ mod tests { use crate::expressions::Column; use crate::projection::ProjectionTargets; + use arrow::array::types::IntervalMonthDayNanoType; 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}; + use datafusion_functions::datetime::date_bin; struct PartitioningTestFixture { schema: SchemaRef, @@ -1097,6 +1424,181 @@ mod tests { ) } + fn timestamp_fixture() -> Result { + PartitioningTestFixture::new(vec![( + "ts", + DataType::Timestamp(TimeUnit::Nanosecond, None), + )]) + } + + fn timestamp_a_fixture() -> Result { + PartitioningTestFixture::new(vec![ + ("ts", DataType::Timestamp(TimeUnit::Nanosecond, None)), + ("a", DataType::Int32), + ]) + } + + fn timestamp_split_point(value: i64) -> SplitPoint { + SplitPoint::new(vec![ScalarValue::TimestampNanosecond(Some(value), None)]) + } + + fn timestamp_a_split_point(timestamp: i64, a: i32) -> SplitPoint { + SplitPoint::new(vec![ + ScalarValue::TimestampNanosecond(Some(timestamp), None), + ScalarValue::Int32(Some(a)), + ]) + } + + fn date_bin_hour_expr( + fixture: &PartitioningTestFixture, + ) -> Result> { + let stride = Arc::new(Literal::new(ScalarValue::IntervalMonthDayNano(Some( + IntervalMonthDayNanoType::make_value(0, 0, 3_600_000_000_000), + )))) as Arc; + Ok(Arc::new(ScalarFunctionExpr::try_new( + date_bin(), + vec![stride, fixture.col(0)], + &fixture.schema, + Arc::new(ConfigOptions::new()), + )?)) + } + + #[test] + fn test_range_partitioning_satisfies_aligned_date_bin() -> Result<()> { + let fixture = timestamp_fixture()?; + let date_bin = date_bin_hour_expr(&fixture)?; + let partitioning = fixture.range_partitioning( + [0], + vec![ + timestamp_split_point(1_704_070_800_000_000_000), + timestamp_split_point(1_704_074_400_000_000_000), + ], + ); + + assert_eq!( + range_partitioning_satisfaction_for_key_partitioning( + &partitioning, + &[Arc::clone(&date_bin)], + &fixture.eq_properties, + false, + ), + PartitioningSatisfaction::Exact + ); + assert_eq!( + range_partitioning_satisfaction_for_key_partitioning( + &partitioning, + &[date_bin], + &fixture.eq_properties, + true, + ), + PartitioningSatisfaction::Exact + ); + + Ok(()) + } + + #[test] + fn test_range_partitioning_does_not_satisfy_crossing_date_bin() -> Result<()> { + let fixture = timestamp_fixture()?; + let date_bin = date_bin_hour_expr(&fixture)?; + let partitioning = fixture.range_partitioning( + [0], + vec![ + timestamp_split_point(1_704_069_000_000_000_000), + timestamp_split_point(1_704_072_600_000_000_000), + ], + ); + + assert_eq!( + range_partitioning_satisfaction_for_key_partitioning( + &partitioning, + &[date_bin], + &fixture.eq_properties, + true, + ), + PartitioningSatisfaction::NotSatisfied + ); + + Ok(()) + } + + #[test] + fn test_hierarchical_range_partitioning_satisfies_aligned_date_bin_subset() + -> Result<()> { + let fixture = timestamp_a_fixture()?; + let date_bin = date_bin_hour_expr(&fixture)?; + let required_exprs = vec![fixture.col(1), Arc::clone(&date_bin)]; + let partitioning = fixture.range_partitioning( + [0, 1], + vec![ + timestamp_a_split_point(1_704_070_800_000_000_000, i32::MIN), + timestamp_a_split_point(1_704_074_400_000_000_000, i32::MIN), + ], + ); + + assert_eq!( + range_partitioning_satisfaction_for_key_partitioning( + &partitioning, + &required_exprs, + &fixture.eq_properties, + true, + ), + PartitioningSatisfaction::Subset + ); + assert_eq!( + range_partitioning_satisfaction_for_key_partitioning( + &partitioning, + &required_exprs, + &fixture.eq_properties, + false, + ), + PartitioningSatisfaction::NotSatisfied + ); + + let mapping = ProjectionMapping::try_new( + vec![ + (fixture.col(1), "a".to_string()), + (date_bin, "bucket".to_string()), + ], + &fixture.schema, + )?; + assert_eq!( + partitioning + .project(&mapping, &fixture.eq_properties) + .to_string(), + "Range([bucket@1 ASC], [(1704070800000000000), (1704074400000000000)], 3)" + ); + + Ok(()) + } + + #[test] + fn test_hierarchical_range_partitioning_rejects_non_boundary_trailing_key() + -> Result<()> { + let fixture = timestamp_a_fixture()?; + let date_bin = date_bin_hour_expr(&fixture)?; + let required_exprs = vec![fixture.col(1), date_bin]; + let partitioning = fixture.range_partitioning( + [0, 1], + vec![ + timestamp_a_split_point(1_704_070_800_000_000_000, 0), + timestamp_a_split_point(1_704_074_400_000_000_000, i32::MIN), + ], + ); + + assert_eq!( + range_partitioning_satisfaction_for_key_partitioning( + &partitioning, + &required_exprs, + &fixture.eq_properties, + true, + ), + PartitioningSatisfaction::NotSatisfied + ); + + Ok(()) + } + fn assert_range_try_new_error( ordering: LexOrdering, split_points: Vec, diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index 1fbf8c6fb78cd..e536e9556f212 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use datafusion_common::Result; use datafusion_physical_expr::{ Distribution, EquivalenceProperties, LexOrdering, LexRequirement, Partitioning, - PhysicalExpr, physical_exprs_equal, + PhysicalExpr, range_partitioning_satisfaction_for_key_partitioning, }; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; @@ -174,45 +174,13 @@ pub(crate) fn range_partitioning_satisfies_key_partitioning( eq_properties: &EquivalenceProperties, allow_subset: bool, ) -> bool { - match partitioning { - Partitioning::Range(range) => { - let partition_exprs = range - .ordering() - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .collect::>(); - - if partition_exprs.is_empty() || required_exprs.is_empty() { - return false; - } - - let eq_group = eq_properties.eq_group(); - let normalized_partition_exprs = partition_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - let normalized_required_exprs = required_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - - if physical_exprs_equal( - &normalized_required_exprs, - &normalized_partition_exprs, - ) { - return true; - } - - allow_subset - && normalized_partition_exprs.len() < normalized_required_exprs.len() - && normalized_partition_exprs.iter().all(|partition_expr| { - normalized_required_exprs - .iter() - .any(|required_expr| partition_expr.eq(required_expr)) - }) - } - _ => false, - } + range_partitioning_satisfaction_for_key_partitioning( + partitioning, + required_exprs, + eq_properties, + allow_subset, + ) + .is_satisfied() } /// Checks whether the given operator is a limit; diff --git a/datafusion/physical-plan/src/distribution_requirements.rs b/datafusion/physical-plan/src/distribution_requirements.rs index 9c7a1336c06a3..e964451672f8d 100644 --- a/datafusion/physical-plan/src/distribution_requirements.rs +++ b/datafusion/physical-plan/src/distribution_requirements.rs @@ -17,12 +17,10 @@ //! Input distribution requirements for physical execution plans. -use std::sync::Arc; - use datafusion_common::{Result, internal_err}; use datafusion_physical_expr::{ Distribution, EquivalenceProperties, Partitioning, PartitioningSatisfaction, - PhysicalExpr, physical_exprs_equal, + range_partitioning_satisfaction_for_key_partitioning, }; use crate::execution_plan::{ExecutionPlan, ExecutionPlanProperties, InvariantLevel}; @@ -392,7 +390,7 @@ impl InputDistributionSatisfaction { return PartitioningSatisfaction::NotSatisfied; }; - range_satisfies_key_partitioning( + range_partitioning_satisfaction_for_key_partitioning( partitioning, required_exprs, eq_properties, @@ -425,54 +423,6 @@ fn validate_child_index( /// generally satisfies [`Distribution::KeyPartitioned`] through /// [`Partitioning::satisfaction`]. /// . -fn range_satisfies_key_partitioning( - partitioning: &Partitioning, - required_exprs: &[Arc], - eq_properties: &EquivalenceProperties, - allow_subset: bool, -) -> PartitioningSatisfaction { - let Partitioning::Range(range) = partitioning else { - return PartitioningSatisfaction::NotSatisfied; - }; - - let partition_exprs = range - .ordering() - .iter() - .map(|sort_expr| Arc::clone(&sort_expr.expr)) - .collect::>(); - - if partition_exprs.is_empty() || required_exprs.is_empty() { - return PartitioningSatisfaction::NotSatisfied; - } - - let eq_group = eq_properties.eq_group(); - let normalized_partition_exprs = partition_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - let normalized_required_exprs = required_exprs - .iter() - .map(|expr| eq_group.normalize_expr(Arc::clone(expr))) - .collect::>(); - - if physical_exprs_equal(&normalized_required_exprs, &normalized_partition_exprs) { - return PartitioningSatisfaction::Exact; - } - - if allow_subset - && normalized_partition_exprs.len() < normalized_required_exprs.len() - && normalized_partition_exprs.iter().all(|partition_expr| { - normalized_required_exprs - .iter() - .any(|required_expr| partition_expr.eq(required_expr)) - }) - { - PartitioningSatisfaction::Subset - } else { - PartitioningSatisfaction::NotSatisfied - } -} - fn compatible_co_partitioning_layout( first: &ChildDistributionRequirement, first_partitioning: &Partitioning, diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs index aa741dded77be..70bf4dec8d803 100644 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -19,7 +19,7 @@ use std::fs::{create_dir_all, remove_dir_all, write}; use std::path::Path; use std::sync::Arc; -use arrow::datatypes::{DataType, Field, Schema}; +use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use datafusion::common::{ScalarValue, SplitPoint}; use datafusion::datasource::file_format::csv::CsvFormat; use datafusion::datasource::listing::{ @@ -93,6 +93,115 @@ pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { ], Some(shifted_output_partitioning), ); + + let time_schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), false), + Field::new("value", DataType::Int32, false), + ])); + let time_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("ts").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(1_704_069_000_000_000_000), + None, + )]), + SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(1_704_072_600_000_000_000), + None, + )]), + SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(1_704_076_200_000_000_000), + None, + )]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_csv_listing_table( + ctx, + "range_partitioned_time", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned_time"), + time_schema, + [ + "2024-01-01T00:15:00,1\n", + "2024-01-01T00:45:00,2\n2024-01-01T01:15:00,4\n", + "2024-01-01T01:45:00,8\n2024-01-01T02:15:00,16\n", + "2024-01-01T02:45:00,32\n", + ], + Some(time_output_partitioning), + ); + + let time_a_schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), false), + Field::new("a", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])); + let time_a_split_point = |ts| { + SplitPoint::new(vec![ + ScalarValue::TimestampNanosecond(Some(ts), None), + ScalarValue::Int32(Some(i32::MIN)), + ]) + }; + let time_a_ordering = || vec![col("ts").sort(true, true), col("a").sort(true, true)]; + + let time_a_crossing_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + time_a_ordering(), + vec![ + time_a_split_point(1_704_069_000_000_000_000), + time_a_split_point(1_704_072_600_000_000_000), + time_a_split_point(1_704_076_200_000_000_000), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_csv_listing_table( + ctx, + "range_partitioned_time_a_crossing", + Path::new(env!("CARGO_MANIFEST_DIR")).join( + "test_files/scratch_range_partitioning/range_partitioned_time_a_crossing", + ), + Arc::clone(&time_a_schema), + [ + "2024-01-01T00:15:00,1,1\n2024-01-01T00:15:00,2,10\n", + "2024-01-01T00:45:00,1,2\n2024-01-01T01:15:00,1,4\n2024-01-01T01:15:00,2,20\n", + "2024-01-01T01:45:00,1,8\n2024-01-01T02:15:00,1,16\n", + "2024-01-01T02:45:00,1,32\n2024-01-01T03:15:00,1,64\n", + ], + Some(time_a_crossing_output_partitioning), + ); + + let time_a_aligned_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + time_a_ordering(), + vec![ + time_a_split_point(1_704_070_800_000_000_000), + time_a_split_point(1_704_074_400_000_000_000), + time_a_split_point(1_704_078_000_000_000_000), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_csv_listing_table( + ctx, + "range_partitioned_time_a_aligned", + Path::new(env!("CARGO_MANIFEST_DIR")).join( + "test_files/scratch_range_partitioning/range_partitioned_time_a_aligned", + ), + time_a_schema, + [ + "2024-01-01T00:15:00,1,1\n2024-01-01T00:45:00,1,2\n2024-01-01T00:15:00,2,10\n", + "2024-01-01T01:15:00,1,4\n2024-01-01T01:45:00,1,8\n2024-01-01T01:15:00,2,20\n", + "2024-01-01T02:15:00,1,16\n2024-01-01T02:45:00,1,32\n", + "2024-01-01T03:15:00,1,64\n", + ], + Some(time_a_aligned_output_partitioning), + ); } fn register_csv_listing_table( diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 5d004ca7fdfa5..844a66a7fdb61 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -22,6 +22,14 @@ # partition 1: range_key in [10, 20), rows (10, 1, 100), (15, 2, 150) # partition 2: range_key in [20, 30), rows (20, 1, 200), (25, 2, 250) # partition 3: range_key in [30, ...), rows (30, 1, 300), (35, 2, 350) +# +# It also registers range_partitioned_time(ts, value), split at half-hour +# boundaries. This makes hourly date_bin buckets cross range partitions. +# +# It also registers range_partitioned_time_a_crossing(ts, a, value) and +# range_partitioned_time_a_aligned(ts, a, value), both range partitioned by +# (ts, a). The first uses half-hour timestamp boundaries and the second uses +# hourly timestamp boundaries. statement ok set datafusion.explain.physical_plan_only = true; @@ -260,6 +268,119 @@ set datafusion.execution.target_partitions = 4; statement ok reset datafusion.optimizer.preserve_file_partitions; +########## +# TEST 8A: Date Bin Aggregate Repartitions Range Partitioned Timestamps +# Range([ts]) cannot satisfy GROUP BY date_bin('1 hour', ts) because the raw +# timestamp range boundaries are not aligned with the hourly bucket boundaries. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT + date_bin(INTERVAL '1 hour', ts) AS bucket, + SUM(value) +FROM range_partitioned_time +GROUP BY date_bin(INTERVAL '1 hour', ts); +---- +physical_plan +01)ProjectionExec: expr=[date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time.ts)@0 as bucket, sum(range_partitioned_time.value)@1 as sum(range_partitioned_time.value)] +02)--AggregateExec: mode=FinalPartitioned, gby=[date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time.ts)@0 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time.ts)], aggr=[sum(range_partitioned_time.value)] +03)----RepartitionExec: partitioning=Hash([date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time.ts)@0], 4), input_partitions=4 +04)------AggregateExec: mode=Partial, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }, ts@0) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time.ts)], aggr=[sum(range_partitioned_time.value)] +05)--------DataSourceExec: file_groups=, projection=[ts, value], output_partitioning=Range([ts@0 ASC], [(1704069000000000000), (1704072600000000000), (1704076200000000000)], 4), file_type=csv, has_header=false + +query PI +SELECT + date_bin(INTERVAL '1 hour', ts) AS bucket, + SUM(value) +FROM range_partitioned_time +GROUP BY date_bin(INTERVAL '1 hour', ts) +ORDER BY bucket; +---- +2024-01-01T00:00:00 3 +2024-01-01T01:00:00 12 +2024-01-01T02:00:00 48 + +########## +# TEST 8B: Hierarchical Date Bin Aggregate Repartitions When Buckets Cross Ranges +# Range([ts, a]) cannot satisfy GROUP BY a, date_bin('1 hour', ts) when +# timestamp range boundaries split hourly buckets. +########## + +query TT +EXPLAIN SELECT + a, + date_bin(INTERVAL '1 hour', ts) AS bucket, + SUM(value) +FROM range_partitioned_time_a_crossing +GROUP BY a, date_bin(INTERVAL '1 hour', ts); +---- +physical_plan +01)ProjectionExec: expr=[a@0 as a, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time_a_crossing.ts)@1 as bucket, sum(range_partitioned_time_a_crossing.value)@2 as sum(range_partitioned_time_a_crossing.value)] +02)--AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time_a_crossing.ts)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time_a_crossing.ts)], aggr=[sum(range_partitioned_time_a_crossing.value)] +03)----RepartitionExec: partitioning=Hash([a@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time_a_crossing.ts)@1], 4), input_partitions=4 +04)------AggregateExec: mode=Partial, gby=[a@1 as a, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }, ts@0) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time_a_crossing.ts)], aggr=[sum(range_partitioned_time_a_crossing.value)] +05)--------DataSourceExec: file_groups=, projection=[ts, a, value], output_partitioning=Range([ts@0 ASC, a@1 ASC], [(1704069000000000000, -2147483648), (1704072600000000000, -2147483648), (1704076200000000000, -2147483648)], 4), file_type=csv, has_header=false + +query IPI +SELECT + a, + date_bin(INTERVAL '1 hour', ts) AS bucket, + SUM(value) +FROM range_partitioned_time_a_crossing +GROUP BY a, date_bin(INTERVAL '1 hour', ts) +ORDER BY a, bucket; +---- +1 2024-01-01T00:00:00 3 +1 2024-01-01T01:00:00 12 +1 2024-01-01T02:00:00 48 +1 2024-01-01T03:00:00 64 +2 2024-01-01T00:00:00 10 +2 2024-01-01T01:00:00 20 + +########## +# TEST 8C: Hierarchical Date Bin Aggregate Reuses Aligned Range Partitioning +# Range([ts, a]) satisfies GROUP BY a, date_bin('1 hour', ts) when every +# timestamp range boundary is an hourly date_bin boundary. +########## + +query TT +EXPLAIN SELECT + a, + date_bin(INTERVAL '1 hour', ts) AS bucket, + SUM(value) +FROM range_partitioned_time_a_aligned +GROUP BY a, date_bin(INTERVAL '1 hour', ts); +---- +physical_plan +01)ProjectionExec: expr=[a@0 as a, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time_a_aligned.ts)@1 as bucket, sum(range_partitioned_time_a_aligned.value)@2 as sum(range_partitioned_time_a_aligned.value)] +02)--AggregateExec: mode=SinglePartitioned, gby=[a@1 as a, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }, ts@0) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 3600000000000 }"),range_partitioned_time_a_aligned.ts)], aggr=[sum(range_partitioned_time_a_aligned.value)] +03)----DataSourceExec: file_groups=, projection=[ts, a, value], output_partitioning=Range([ts@0 ASC, a@1 ASC], [(1704070800000000000, -2147483648), (1704074400000000000, -2147483648), (1704078000000000000, -2147483648)], 4), file_type=csv, has_header=false + +query IPI +SELECT + a, + date_bin(INTERVAL '1 hour', ts) AS bucket, + SUM(value) +FROM range_partitioned_time_a_aligned +GROUP BY a, date_bin(INTERVAL '1 hour', ts) +ORDER BY a, bucket; +---- +1 2024-01-01T00:00:00 3 +1 2024-01-01T01:00:00 12 +1 2024-01-01T02:00:00 48 +1 2024-01-01T03:00:00 64 +2 2024-01-01T00:00:00 10 +2 2024-01-01T01:00:00 20 + statement ok set datafusion.optimizer.prefer_hash_join = true;