jayshrivastava commented on code in PR #24501:
URL: https://github.com/apache/datafusion/pull/24501#discussion_r3831117770


##########
datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt:
##########
@@ -25,19 +25,16 @@
 #   WHERE col4 = 'a'
 #   GROUP BY key, time_bin
 #
-# Scan metadata already advertises:
+# Scan metadata advertises:
 # 1. Range([timestamp]) and output_ordering=[key, timestamp]
 # 2. Two file_groups, so the two 60-minute streams run in parallel
 #
-# Improvement opportunity:
-# date_bin(60s) is monotonic in timestamp and the hour split is aligned to bin
-# boundaries, so (key, time_bin) is partition-disjoint. Aggregation could be a
-# single streaming SinglePartitioned step with no hash shuffle.
+# date_bin(60s) and date_trunc('hour') are monotonic in timestamp and the hour
+# split is aligned to those bins, so (key, time_bin) is partition-disjoint.
+# Aggregation is one streaming SinglePartitioned step with no hash shuffle.
 #
-# Today's plan still hash-repartitions:
-#   Partial AggregateExec (ordering_mode=Sorted)
-#     -> RepartitionExec Hash([key, date_bin(...)])
-#     -> FinalPartitioned AggregateExec (ordering_mode=Sorted)
+# date_trunc('day') bins straddle the hour split, so that query still
+# hash-repartitions.

Review Comment:
   nit: can you use actual numbers as an example to show a date_bin/date_trunc 
which preserves partitioning and one that doesn't?



##########
datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt:
##########
@@ -146,6 +137,64 @@ k2 2024-01-01T00:30:00 7
 k2 2024-01-01T01:30:00 30
 k2 2024-01-01T01:45:00 5
 
+##########
+# TEST 4: date_trunc('hour') is aligned to the hour split, so the same
+# SinglePartitioned streaming plan applies.
+##########
+
+query TT
+EXPLAIN SELECT key, date_trunc('hour', timestamp) AS time_bin, sum(value)
+FROM range_sorted_time_bin
+WHERE col4 = 'a'
+GROUP BY key, time_bin;
+----
+physical_plan
+01)ProjectionExec: expr=[key@0 as key, 
date_trunc(Utf8("hour"),range_sorted_time_bin.timestamp)@1 as time_bin, 
sum(range_sorted_time_bin.value)@2 as sum(range_sorted_time_bin.value)]
+02)--AggregateExec: mode=SinglePartitioned, gby=[key@0 as key, 
date_trunc(hour, timestamp@1) as 
date_trunc(Utf8("hour"),range_sorted_time_bin.timestamp)], 
aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted
+03)----FilterExec: col4@1 = a, projection=[key@0, timestamp@2, value@3]
+04)------DataSourceExec: file_groups={2 groups: 
[[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-0.parquet],
 
[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-1.parquet]]},
 projection=[key, col4, timestamp, value], output_ordering=[key@0 ASC, 
timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], 
[(1704070800000000000)], 2), file_type=parquet, predicate=col4@4 = a, 
pruning_predicate=col4_null_count@2 != row_count@3 AND col4_min@0 <= a AND a <= 
col4_max@1, required_guarantees=[col4 in (a)]
+
+query TPI
+SELECT key, date_trunc('hour', timestamp) AS time_bin, sum(value)
+FROM range_sorted_time_bin
+WHERE col4 = 'a'
+GROUP BY key, time_bin
+ORDER BY key, time_bin;
+----
+k1 2024-01-01T00:00:00 3
+k1 2024-01-01T01:00:00 30
+k2 2024-01-01T00:00:00 7
+k2 2024-01-01T01:00:00 35
+
+##########
+# TEST 5: date_trunc('day') bins straddle the hour split (both file groups are
+# 2024-01-01), so grouping is not partition-disjoint and a hash shuffle 
remains.
+##########
+
+query TT
+EXPLAIN SELECT key, date_trunc('day', timestamp) AS time_bin, sum(value)
+FROM range_sorted_time_bin
+WHERE col4 = 'a'
+GROUP BY key, time_bin;
+----
+physical_plan
+01)ProjectionExec: expr=[key@0 as key, 
date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1 as time_bin, 
sum(range_sorted_time_bin.value)@2 as sum(range_sorted_time_bin.value)]
+02)--AggregateExec: mode=FinalPartitioned, gby=[key@0 as key, 
date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1 as 
date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)], 
aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted
+03)----RepartitionExec: partitioning=Hash([key@0, 
date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1], 2), 
input_partitions=2, preserve_order=true, sort_exprs=key@0 ASC, 
date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)@1 ASC
+04)------AggregateExec: mode=Partial, gby=[key@0 as key, date_trunc(day, 
timestamp@1) as date_trunc(Utf8("day"),range_sorted_time_bin.timestamp)], 
aggr=[sum(range_sorted_time_bin.value)], ordering_mode=Sorted
+05)--------FilterExec: col4@1 = a, projection=[key@0, timestamp@2, value@3]
+06)----------DataSourceExec: file_groups={2 groups: 
[[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-0.parquet],
 
[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_sorted_time_bin/part-1.parquet]]},
 projection=[key, col4, timestamp, value], output_ordering=[key@0 ASC, 
timestamp@2 ASC], output_partitioning=Range([timestamp@2 ASC], 
[(1704070800000000000)], 2), file_type=parquet, predicate=col4@4 = a, 
pruning_predicate=col4_null_count@2 != row_count@3 AND col4_min@0 <= a AND a <= 
col4_max@1, required_guarantees=[col4 in (a)]
+
+query TPI
+SELECT key, date_trunc('day', timestamp) AS time_bin, sum(value)
+FROM range_sorted_time_bin
+WHERE col4 = 'a'
+GROUP BY key, time_bin
+ORDER BY key, time_bin;
+----
+k1 2024-01-01T00:00:00 33
+k2 2024-01-01T00:00:00 42

Review Comment:
   These tests look good 👍🏽 



##########
datafusion/physical-expr/src/equivalence/properties/mod.rs:
##########
@@ -1319,6 +1319,31 @@ impl EquivalenceProperties {
             .unwrap_or_else(|_| ExprProperties::new_unknown())
     }
 
+    /// Returns true when `expr` is a (possibly non-strict) monotonic function 
of
+    /// `range_key` plus literals, such as `date_bin(interval, timestamp)` or
+    /// `date_trunc(unit, timestamp)`.
+    ///
+    /// The identity `expr == range_key` returns false so callers can treat 
"emit
+    /// the key as-is" separately from "emit a function of the key".
+    pub(crate) fn is_monotonic_function_of(
+        &self,
+        expr: &Arc<dyn PhysicalExpr>,
+        range_key: &Arc<dyn PhysicalExpr>,
+    ) -> bool {
+        if expr.eq(range_key) {
+            return false;
+        }
+        let dependencies = 
Dependencies::new(std::iter::once(PhysicalSortExpr::new(
+            Arc::clone(range_key),
+            Default::default(),
+        )));
+        matches!(
+            get_expr_properties(expr, &dependencies, &self.schema)
+                .map(|properties| properties.sort_properties),
+            Ok(SortProperties::Ordered(_))

Review Comment:
   This doesn't capture the ordering, so a non monotonic function like `f(x) = 
-x` will return true here even though it flips the ordering from ASC to DESC. 
Even if `-x` works for range partitioning, this function may be used for some 
different use case, in which it will incorrectly return `true`.
   
    Let's just assert that the options in `SortProperties::Ordered(options)` 
are the same as `sort_expr.options`  and add a test for `-x`.
   



##########
datafusion/physical-expr/src/equivalence/properties/mod.rs:
##########
@@ -1319,6 +1319,31 @@ impl EquivalenceProperties {
             .unwrap_or_else(|_| ExprProperties::new_unknown())
     }
 
+    /// Returns true when `expr` is a (possibly non-strict) monotonic function 
of
+    /// `range_key` plus literals, such as `date_bin(interval, timestamp)` or
+    /// `date_trunc(unit, timestamp)`.
+    ///
+    /// The identity `expr == range_key` returns false so callers can treat 
"emit
+    /// the key as-is" separately from "emit a function of the key".
+    pub(crate) fn is_monotonic_function_of(

Review Comment:
   I'd rename this.
   
   - Doing `equivalence_properties.is_monotonic_function_of` sounds like the 
properies are a monotonic function.
   - The two expressions aren't necessarily either functions. They are related 
via the `Dependencies` relationship.
   
   Maybe rename to `check_monotonic_dependency` or `check_monotonic_transform`?



##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -433,12 +570,24 @@ impl Partitioning {
                         .iter()
                         .map(|sort_expr| Arc::clone(&sort_expr.expr))
                         .collect::<Vec<_>>();
-                    Self::key_satisfaction(
+                    let satisfaction = Self::key_satisfaction(
                         &partition_exprs,
                         required_exprs,
                         eq_properties,
                         allow_subset,
-                    )
+                    );
+                    if satisfaction == PartitioningSatisfaction::NotSatisfied
+                        && allow_subset

Review Comment:
   If partitioned on just `timestamp` and you do `select date_bin(timestamp) 
from foo`, then we should check `range_monotonic_fn_satisfies_keys` even if 
`allow_subset` is false.



##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -252,37 +258,168 @@ impl RangePartitioning {
     ///
     /// Returns `None` if any range key cannot be projected or if projection
     /// collapses distinct range keys into duplicate output expressions.
+    ///
+    /// If a projection drops a range key but keeps a monotonic function of it
+    /// (for example `date_bin(interval, timestamp)` or `date_trunc(unit, 
timestamp)`
+    /// while range-partitioned on `timestamp`), the range can still be 
projected.
+    /// Adjacent partitions stay disjoint only when evaluating the function at
+    /// each split point and its predecessor yields different values, so bins
+    /// do not straddle file groups.
     fn project(
         &self,
         mapping: &ProjectionMapping,
         input_eq_properties: &EquivalenceProperties,
     ) -> Option<Self> {
-        let exprs = self
-            .ordering
-            .iter()
-            .map(|sort_expr| Arc::clone(&sort_expr.expr))
-            .collect::<Vec<_>>();
-        let projected_exprs = input_eq_properties
-            .project_expressions(&exprs, mapping)
-            .collect::<Option<Vec<_>>>()?;
-        let sort_exprs = self
-            .ordering
-            .iter()
-            .zip(projected_exprs)
-            .map(|(sort_expr, expr)| PhysicalSortExpr::new(expr, 
sort_expr.options))
-            .collect::<Vec<_>>();
+        let mut split_points = self.split_points.clone();
+        let mut sort_exprs = Vec::with_capacity(self.ordering.len());
+        for (key_idx, sort_expr) in self.ordering.iter().enumerate() {
+            if let Some(projected) =
+                input_eq_properties.project_expr(&sort_expr.expr, mapping)
+            {
+                sort_exprs.push(PhysicalSortExpr::new(projected, 
sort_expr.options));
+                continue;
+            }
+
+            let (target, source) =
+                monotonic_range_key_projection(sort_expr, mapping, 
input_eq_properties)?;
+            if !monotonic_fn_keeps_partitions_disjoint(
+                &source,
+                &sort_expr.expr,
+                &split_points,
+                key_idx,
+            ) {
+                return None;
+            }
+            if let Some(updated) = project_split_points_through_fn(
+                &source,
+                &sort_expr.expr,
+                &split_points,
+                key_idx,
+            ) {
+                split_points = updated;
+            }
+            sort_exprs.push(PhysicalSortExpr::new(target, sort_expr.options));
+        }
         let ordering = LexOrdering::new(sort_exprs)?;
         if ordering.len() != self.ordering.len() {
             return None;
         }
 
         Some(Self {
             ordering,
-            split_points: self.split_points.clone(),
+            split_points,
         })
     }
 }
 
+/// Finds a projection mapping whose source is a monotonic function of 
`sort_expr`.
+fn monotonic_range_key_projection(
+    sort_expr: &PhysicalSortExpr,
+    mapping: &ProjectionMapping,
+    eq_properties: &EquivalenceProperties,
+) -> Option<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)> {
+    mapping.iter().find_map(|(source, targets)| {
+        eq_properties
+            .is_monotonic_function_of(source, &sort_expr.expr)
+            .then(|| (Arc::clone(&targets.first().0), Arc::clone(source)))
+    })
+}
+
+/// Adjacent range partitions remain disjoint on `fn_expr` when the function
+/// value at each split differs from the value immediately below the split.
+fn monotonic_fn_keeps_partitions_disjoint(
+    fn_expr: &Arc<dyn PhysicalExpr>,
+    range_key: &Arc<dyn PhysicalExpr>,
+    split_points: &[SplitPoint],
+    key_idx: usize,
+) -> bool {
+    split_points.iter().all(|split_point| {
+        let Some(split_value) = split_point.values().get(key_idx) else {
+            return false;
+        };
+        let Some(predecessor) = checked_predecessor(split_value) else {
+            return false;
+        };
+        let Some(at_split) = evaluate_expr_on_key(fn_expr, range_key, 
split_value) else {
+            return false;
+        };
+        let Some(below_split) = evaluate_expr_on_key(fn_expr, range_key, 
&predecessor)
+        else {
+            return false;
+        };
+        at_split != below_split
+    })
+}
+
+fn project_split_points_through_fn(
+    fn_expr: &Arc<dyn PhysicalExpr>,
+    range_key: &Arc<dyn PhysicalExpr>,
+    split_points: &[SplitPoint],
+    key_idx: usize,
+) -> Option<Vec<SplitPoint>> {
+    split_points
+        .iter()
+        .map(|split_point| {
+            let split_value = split_point.values().get(key_idx)?;
+            let projected = evaluate_expr_on_key(fn_expr, range_key, 
split_value)?;
+            let mut values = split_point.values().to_vec();
+            values[key_idx] = projected;
+            Some(SplitPoint::new(values))
+        })
+        .collect()
+}
+
+/// Evaluates `expr` after substituting `range_key` with `value`.
+fn evaluate_expr_on_key(
+    expr: &Arc<dyn PhysicalExpr>,
+    range_key: &Arc<dyn PhysicalExpr>,
+    value: &ScalarValue,
+) -> Option<ScalarValue> {
+    let literal: Arc<dyn PhysicalExpr> = Arc::new(Literal::new(value.clone()));
+    let rewritten = Arc::clone(expr)
+        .transform(|node| {
+            if node.eq(range_key) {
+                Ok(Transformed::yes(Arc::clone(&literal)))
+            } else {
+                Ok(Transformed::no(node))
+            }
+        })
+        .ok()?;
+    if !rewritten.transformed {
+        return None;
+    }
+    let batch = create_dummy_batch().ok()?;
+    match rewritten.data.evaluate(batch).ok()? {
+        ColumnarValue::Scalar(scalar) => Some(scalar),
+        ColumnarValue::Array(array) => ScalarValue::try_from_array(&array, 
0).ok(),
+    }
+}
+
+/// `Range([x])` satisfies grouping by `(..., f(x), ...)` when `f` is 
monotonic in
+/// `x` and adjacent partitions do not share `f` values (bins do not straddle
+/// split points). That makes `(key, date_bin(timestamp))` and
+/// `(key, date_trunc(timestamp))` partition-disjoint when the table is
+/// range-partitioned on `timestamp` and the split is aligned to the bin.
+fn range_monotonic_fn_satisfies_keys(
+    range: &RangePartitioning,
+    required_exprs: &[Arc<dyn PhysicalExpr>],
+    eq_properties: &EquivalenceProperties,
+) -> bool {
+    if range.ordering().len() != 1 {

Review Comment:
   Do you plan to expand this in a follow up?



##########
datafusion/physical-expr/src/equivalence/properties/mod.rs:
##########
@@ -1319,6 +1319,31 @@ impl EquivalenceProperties {
             .unwrap_or_else(|_| ExprProperties::new_unknown())
     }
 
+    /// Returns true when `expr` is a (possibly non-strict) monotonic function 
of
+    /// `range_key` plus literals, such as `date_bin(interval, timestamp)` or
+    /// `date_trunc(unit, timestamp)`.
+    ///
+    /// The identity `expr == range_key` returns false so callers can treat 
"emit
+    /// the key as-is" separately from "emit a function of the key".
+    pub(crate) fn is_monotonic_function_of(
+        &self,
+        expr: &Arc<dyn PhysicalExpr>,
+        range_key: &Arc<dyn PhysicalExpr>,

Review Comment:
   Rename these to source_expr and transformed_expr? Definately would not call 
this `range_key` because this function has nothing to do with range 
partitioning. 



##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -252,37 +258,168 @@ impl RangePartitioning {
     ///
     /// Returns `None` if any range key cannot be projected or if projection
     /// collapses distinct range keys into duplicate output expressions.
+    ///
+    /// If a projection drops a range key but keeps a monotonic function of it
+    /// (for example `date_bin(interval, timestamp)` or `date_trunc(unit, 
timestamp)`
+    /// while range-partitioned on `timestamp`), the range can still be 
projected.
+    /// Adjacent partitions stay disjoint only when evaluating the function at
+    /// each split point and its predecessor yields different values, so bins
+    /// do not straddle file groups.
     fn project(
         &self,
         mapping: &ProjectionMapping,
         input_eq_properties: &EquivalenceProperties,
     ) -> Option<Self> {
-        let exprs = self
-            .ordering
-            .iter()
-            .map(|sort_expr| Arc::clone(&sort_expr.expr))
-            .collect::<Vec<_>>();
-        let projected_exprs = input_eq_properties
-            .project_expressions(&exprs, mapping)
-            .collect::<Option<Vec<_>>>()?;
-        let sort_exprs = self
-            .ordering
-            .iter()
-            .zip(projected_exprs)
-            .map(|(sort_expr, expr)| PhysicalSortExpr::new(expr, 
sort_expr.options))
-            .collect::<Vec<_>>();
+        let mut split_points = self.split_points.clone();
+        let mut sort_exprs = Vec::with_capacity(self.ordering.len());
+        for (key_idx, sort_expr) in self.ordering.iter().enumerate() {
+            if let Some(projected) =
+                input_eq_properties.project_expr(&sort_expr.expr, mapping)

Review Comment:
   Can you add a test where we're range partitioned on `a` and project `select 
a, date_bin(a) from foo`? 
   
   I think this case will preserve range partitioning on `a`. This behavior is 
a bit subtle. Your comment is good though "If a projection drops a range key 
but keeps a monotonic function of it".
   
   Testing `select a as b, date_bin(a) as bucket from foo` would be good too. 
This should preserve the same range partitioning.



##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -252,37 +258,168 @@ impl RangePartitioning {
     ///
     /// Returns `None` if any range key cannot be projected or if projection
     /// collapses distinct range keys into duplicate output expressions.
+    ///
+    /// If a projection drops a range key but keeps a monotonic function of it
+    /// (for example `date_bin(interval, timestamp)` or `date_trunc(unit, 
timestamp)`
+    /// while range-partitioned on `timestamp`), the range can still be 
projected.
+    /// Adjacent partitions stay disjoint only when evaluating the function at
+    /// each split point and its predecessor yields different values, so bins
+    /// do not straddle file groups.
     fn project(
         &self,
         mapping: &ProjectionMapping,
         input_eq_properties: &EquivalenceProperties,
     ) -> Option<Self> {
-        let exprs = self
-            .ordering
-            .iter()
-            .map(|sort_expr| Arc::clone(&sort_expr.expr))
-            .collect::<Vec<_>>();
-        let projected_exprs = input_eq_properties
-            .project_expressions(&exprs, mapping)
-            .collect::<Option<Vec<_>>>()?;
-        let sort_exprs = self
-            .ordering
-            .iter()
-            .zip(projected_exprs)
-            .map(|(sort_expr, expr)| PhysicalSortExpr::new(expr, 
sort_expr.options))
-            .collect::<Vec<_>>();
+        let mut split_points = self.split_points.clone();
+        let mut sort_exprs = Vec::with_capacity(self.ordering.len());
+        for (key_idx, sort_expr) in self.ordering.iter().enumerate() {
+            if let Some(projected) =
+                input_eq_properties.project_expr(&sort_expr.expr, mapping)
+            {
+                sort_exprs.push(PhysicalSortExpr::new(projected, 
sort_expr.options));
+                continue;
+            }
+
+            let (target, source) =
+                monotonic_range_key_projection(sort_expr, mapping, 
input_eq_properties)?;
+            if !monotonic_fn_keeps_partitions_disjoint(
+                &source,
+                &sort_expr.expr,
+                &split_points,
+                key_idx,
+            ) {
+                return None;
+            }
+            if let Some(updated) = project_split_points_through_fn(
+                &source,
+                &sort_expr.expr,
+                &split_points,
+                key_idx,
+            ) {
+                split_points = updated;
+            }
+            sort_exprs.push(PhysicalSortExpr::new(target, sort_expr.options));
+        }
         let ordering = LexOrdering::new(sort_exprs)?;
         if ordering.len() != self.ordering.len() {
             return None;
         }
 
         Some(Self {
             ordering,
-            split_points: self.split_points.clone(),
+            split_points,
         })
     }
 }
 
+/// Finds a projection mapping whose source is a monotonic function of 
`sort_expr`.
+fn monotonic_range_key_projection(
+    sort_expr: &PhysicalSortExpr,
+    mapping: &ProjectionMapping,
+    eq_properties: &EquivalenceProperties,
+) -> Option<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)> {
+    mapping.iter().find_map(|(source, targets)| {
+        eq_properties
+            .is_monotonic_function_of(source, &sort_expr.expr)
+            .then(|| (Arc::clone(&targets.first().0), Arc::clone(source)))
+    })
+}
+
+/// Adjacent range partitions remain disjoint on `fn_expr` when the function
+/// value at each split differs from the value immediately below the split.
+fn monotonic_fn_keeps_partitions_disjoint(
+    fn_expr: &Arc<dyn PhysicalExpr>,
+    range_key: &Arc<dyn PhysicalExpr>,
+    split_points: &[SplitPoint],
+    key_idx: usize,
+) -> bool {
+    split_points.iter().all(|split_point| {

Review Comment:
   This is pretty clean. My approach would have been to store the previous 
mapped split point in a variable and check monotonicity that way



##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -252,37 +258,168 @@ impl RangePartitioning {
     ///
     /// Returns `None` if any range key cannot be projected or if projection
     /// collapses distinct range keys into duplicate output expressions.
+    ///
+    /// If a projection drops a range key but keeps a monotonic function of it
+    /// (for example `date_bin(interval, timestamp)` or `date_trunc(unit, 
timestamp)`
+    /// while range-partitioned on `timestamp`), the range can still be 
projected.
+    /// Adjacent partitions stay disjoint only when evaluating the function at
+    /// each split point and its predecessor yields different values, so bins
+    /// do not straddle file groups.
     fn project(
         &self,
         mapping: &ProjectionMapping,
         input_eq_properties: &EquivalenceProperties,
     ) -> Option<Self> {
-        let exprs = self
-            .ordering
-            .iter()
-            .map(|sort_expr| Arc::clone(&sort_expr.expr))
-            .collect::<Vec<_>>();
-        let projected_exprs = input_eq_properties
-            .project_expressions(&exprs, mapping)
-            .collect::<Option<Vec<_>>>()?;
-        let sort_exprs = self
-            .ordering
-            .iter()
-            .zip(projected_exprs)
-            .map(|(sort_expr, expr)| PhysicalSortExpr::new(expr, 
sort_expr.options))
-            .collect::<Vec<_>>();
+        let mut split_points = self.split_points.clone();
+        let mut sort_exprs = Vec::with_capacity(self.ordering.len());
+        for (key_idx, sort_expr) in self.ordering.iter().enumerate() {
+            if let Some(projected) =
+                input_eq_properties.project_expr(&sort_expr.expr, mapping)
+            {
+                sort_exprs.push(PhysicalSortExpr::new(projected, 
sort_expr.options));
+                continue;
+            }
+
+            let (target, source) =
+                monotonic_range_key_projection(sort_expr, mapping, 
input_eq_properties)?;
+            if !monotonic_fn_keeps_partitions_disjoint(
+                &source,
+                &sort_expr.expr,
+                &split_points,
+                key_idx,
+            ) {
+                return None;
+            }
+            if let Some(updated) = project_split_points_through_fn(
+                &source,
+                &sort_expr.expr,
+                &split_points,
+                key_idx,
+            ) {
+                split_points = updated;
+            }
+            sort_exprs.push(PhysicalSortExpr::new(target, sort_expr.options));
+        }
         let ordering = LexOrdering::new(sort_exprs)?;
         if ordering.len() != self.ordering.len() {
             return None;
         }
 
         Some(Self {
             ordering,
-            split_points: self.split_points.clone(),
+            split_points,
         })
     }
 }
 
+/// Finds a projection mapping whose source is a monotonic function of 
`sort_expr`.
+fn monotonic_range_key_projection(
+    sort_expr: &PhysicalSortExpr,
+    mapping: &ProjectionMapping,
+    eq_properties: &EquivalenceProperties,
+) -> Option<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)> {
+    mapping.iter().find_map(|(source, targets)| {
+        eq_properties
+            .is_monotonic_function_of(source, &sort_expr.expr)
+            .then(|| (Arc::clone(&targets.first().0), Arc::clone(source)))

Review Comment:
   I think it's okay to only look at the first target. Target just contains the 
alias IIUC. Ex.
   `select non_monotonic_function(a), date_bin(a) from foo` has 2 sources 
(which contain the projection) with 1 target each, so this function works and 
should find the monotonic date_bin function. Can you add a test for this query 
which calls this function and ensure that it correctly finds the first 
monotonic function?
   
   `select non_monotonic_function(a), date_bin(a) as bin1, date_bin(a) as bin2 
from foo` would also be interesting.
   
   Also note that the way the code is written, `select date_bin(a, 45s) as 
bin1, date_bin(a, 60s) as bin2 from foo` will only check the **first** date_bin.
   
   What if the 45s straddles the partitions but 60s doesn't? Then the range 
partitioning will fall back to `Unknown` right? Maybe we should search for ALL 
the monotonic functions here and then in `RangePartitioning::project(...)`, we 
check each one to see `monotonic_fn_keeps_partitions_disjoint` and use the 
first one that meets both requirements.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to