consider date_bin in range partition satisfaction

datafusion/expr/src/lib.rs

@@ -135,7 +135,8 @@ pub use udaf::{
     udaf_default_window_function_schema_name,
 };
 pub use udf::{
-    ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, StructFieldMapping,
+    RangePartitioningTransform, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF,
+    ScalarUDFImpl, StructFieldMapping,
 };
 pub use udwf::{LimitEffect, ReversedUDWF, WindowUDF, WindowUDFImpl};
 pub use window_frame::{WindowFrame, WindowFrameBound, WindowFrameUnits};

datafusion/expr/src/udf.rs

@@ -57,6 +57,32 @@ pub struct StructFieldMapping {
     pub fields: Vec<(Vec<ScalarValue>, usize)>,
 }
 
+/// Describes how a scalar UDF can transform range partitioning boundaries.
+///
+/// For example, data may be range partitioned on column `timestamp`
+/// at split points 600, 1200, 1800 etc. Applying a projection such as
+/// `date_bin(timestamp, 30 seconds) as date_bin_timestamp` should
+/// preserve range partitioning on the projected column `date_bin_timestamp`.
+///
+/// Implementations of [`RangePartitioningTransform`]
+/// need to do the following:
+/// (a) identify the function argument which is being transformed
+/// (b) map values from input to the UDF output
+///
+/// The physical partitioning machinery the remaining work and safety checks.
+/// For example, the mapped range partition boundaries must be sorted, distinct,
+/// and non null.
+pub trait RangePartitioningTransform: Debug + Send + Sync {
+    /// Index of the function argument which is being mapped.
+    fn source_index(&self) -> usize;
+
+    /// Maps a source value to the corresponding output value. Will be called
+    /// using range partition split points.
+    ///
+    /// Returns `None` when the boundary value cannot be mapped safely.
+    fn map_boundary_value(&self, value: &ScalarValue) -> Option<ScalarValue>;
+}
+
 /// Logical representation of a Scalar User Defined Function.
 ///
 /// A scalar function produces a single row output for each row of input. This
@@ -340,6 +366,14 @@ impl ScalarUDF {
         self.inner.struct_field_mapping(literal_args)
     }
 
+    /// See [`ScalarUDFImpl::range_partitioning_transform`] for more details.
+    pub fn range_partitioning_transform(
+        &self,
+        literal_args: &[Option<ScalarValue>],
+    ) -> Option<Box<dyn RangePartitioningTransform>> {
+        self.inner.range_partitioning_transform(literal_args)
+    }
+
     /// Updates bounds for child expressions, given a known interval for this
     /// function. This is used to propagate constraints down through an expression
     /// tree.
@@ -1029,6 +1063,20 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any {
         None
     }
 
+    /// Returns a transform describing how this function maps range
+    /// partitioning boundaries.
+    ///
+    /// `literal_args[i]` is `Some(value)` when argument `i` is a known literal.
+    /// Implementations should return a transform only when those literals make
+    /// the function a monotonic, many-to-one transformation of one input and
+    /// mapping an aligned boundary preserves partition ownership.
+    fn range_partitioning_transform(
+        &self,
+        _literal_args: &[Option<ScalarValue>],
+    ) -> Option<Box<dyn RangePartitioningTransform>> {
+        None
+    }
+
     /// Returns the documentation for this Scalar UDF.
     ///
     /// Documentation can be accessed programmatically as well as generating
@@ -1188,6 +1236,13 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl {
         self.inner.struct_field_mapping(literal_args)
     }
 
+    fn range_partitioning_transform(
+        &self,
+        literal_args: &[Option<ScalarValue>],
+    ) -> Option<Box<dyn RangePartitioningTransform>> {
+        self.inner.range_partitioning_transform(literal_args)
+    }
+
     fn output_ordering(&self, inputs: &[ExprProperties]) -> Result<SortProperties> {
         self.inner.output_ordering(inputs)
     }

datafusion/functions/src/datetime/date_bin.rs

@@ -40,8 +40,8 @@ use datafusion_common::{
 use datafusion_expr::TypeSignature::Exact;
 use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
 use datafusion_expr::{
-    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
-    TIMEZONE_WILDCARD, Volatility,
+    ColumnarValue, Documentation, RangePartitioningTransform, ScalarFunctionArgs,
+    ScalarUDFImpl, Signature, TIMEZONE_WILDCARD, Volatility,
 };
 use datafusion_macros::user_doc;
 
@@ -117,6 +117,32 @@ pub struct DateBinFunc {
     signature: Signature,
 }
 
+#[derive(Debug)]
+struct DateBinRangePartitioningTransform {
+    stride: ScalarValue,
+    origin: ScalarValue,
+}
+
+impl RangePartitioningTransform for DateBinRangePartitioningTransform {
+    fn source_index(&self) -> usize {
+        1
+    }
+
+    fn map_boundary_value(&self, value: &ScalarValue) -> Option<ScalarValue> {
+        let result = date_bin_impl(
+            &ColumnarValue::Scalar(self.stride.clone()),
+            &ColumnarValue::Scalar(value.clone()),
+            &ColumnarValue::Scalar(self.origin.clone()),
+        )
+        .ok()?;
+
+        match result {
+            ColumnarValue::Scalar(value) => Some(value),
+            ColumnarValue::Array(_) => None,
+        }
+    }
+}
+
 impl Default for DateBinFunc {
     fn default() -> Self {
         Self::new()
@@ -283,11 +309,61 @@ impl ScalarUDFImpl for DateBinFunc {
             Ok(SortProperties::Unordered)
         }
     }
+
+    fn range_partitioning_transform(
+        &self,
+        literal_args: &[Option<ScalarValue>],
+    ) -> Option<Box<dyn RangePartitioningTransform>> {
+        if !(2..=3).contains(&literal_args.len()) {
+            return None;
+        }
+
+        let stride = literal_args.first()?.as_ref()?;
+        if !is_positive_fixed_width_stride(stride) {
+            return None;
+        }
+
+        let origin = if literal_args.len() == 3 {
+            let origin = literal_args.get(2)?.as_ref()?;
+            if !matches!(origin, ScalarValue::TimestampNanosecond(Some(_), _)) {
+                return None;
+            }
+            origin.clone()
+        } else {
+            ScalarValue::TimestampNanosecond(Some(0), Some("+00:00".into()))
+        };
+
+        Some(Box::new(DateBinRangePartitioningTransform {
+            stride: stride.clone(),
+            origin,
+        }))
+    }
+
     fn documentation(&self) -> Option<&Documentation> {
         self.doc()
     }
 }
 
+fn is_positive_fixed_width_stride(value: &ScalarValue) -> bool {
+    let nanos = match value {
+        ScalarValue::IntervalDayTime(Some(value)) => {
+            let (days, millis) = IntervalDayTimeType::to_parts(*value);
+            i128::from(days) * i128::from(NANOSECONDS_IN_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 false;
+            }
+            i128::from(days) * i128::from(NANOSECONDS_IN_DAY) + i128::from(nanos)
+        }
+        _ => return false,
+    };
+
+    nanos > 0 && nanos <= i128::from(i64::MAX)
+}
+
 const NANOS_PER_MICRO: i64 = 1_000;
 const NANOS_PER_MILLI: i64 = 1_000_000;
 const NANOS_PER_SEC: i64 = NANOSECONDS;
@@ -851,6 +927,62 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_range_partitioning_transform() {
+        let hour = ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano {
+            months: 0,
+            days: 0,
+            nanoseconds: 3_600_000_000_000,
+        }));
+        let transform = DateBinFunc::new()
+            .range_partitioning_transform(&[Some(hour), None])
+            .expect("fixed-width date_bin should transform range boundaries");
+
+        assert_eq!(transform.source_index(), 1);
+        assert_eq!(
+            transform.map_boundary_value(&ScalarValue::TimestampNanosecond(
+                Some(3_600_000_000_000),
+                None,
+            )),
+            Some(ScalarValue::TimestampNanosecond(
+                Some(3_600_000_000_000),
+                None,
+            ))
+        );
+        assert_eq!(
+            transform.map_boundary_value(&ScalarValue::TimestampNanosecond(
+                Some(5_400_000_000_000),
+                None,
+            )),
+            Some(ScalarValue::TimestampNanosecond(
+                Some(3_600_000_000_000),
+                None,
+            ))
+        );
+
+        let month = ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano {
+            months: 1,
+            days: 0,
+            nanoseconds: 0,
+        }));
+        assert!(
+            DateBinFunc::new()
+                .range_partitioning_transform(&[Some(month), None])
+                .is_none()
+        );
+
+        let negative = ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNano {
+            months: 0,
+            days: 0,
+            nanoseconds: -1,
+        }));
+        assert!(
+            DateBinFunc::new()
+                .range_partitioning_transform(&[Some(negative), None])
+                .is_none()
+        );
+    }
+
     #[test]
     fn test_date_bin() {
         let return_field = &Arc::new(Field::new(

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,