jayzhan211 commented on code in PR #24501:
URL: https://github.com/apache/datafusion/pull/24501#discussion_r3938466102
##########
datafusion/physical-expr/src/partitioning.rs:
##########
@@ -252,34 +258,215 @@ 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
+ /// when that key is the **final** range key. Flooring transforms are not
+ /// injective: rows above a split can collapse onto `f(split)` and then be
+ /// reordered by trailing keys, so the same fallback on an earlier key
would
+ /// move rows across the projected boundary.
+ /// Adjacent partitions stay disjoint only when evaluating the function at
+ /// each split point and the adjacent value on the open side of the split
+ /// (predecessor for ASC, successor for DESC) 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());
+ let last_key_idx = self.ordering.len().checked_sub(1)?;
+ 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;
+ }
+
+ // A non-injective transform preserves the partition boundary
+ // only for the final range key. For an earlier key, rows above
+ // the split can collapse onto `f(split)` and are then ordered
+ // by the trailing keys, crossing the projected boundary.
+ if key_idx != last_key_idx {
+ return None;
+ }
+
+ let (target, source) =
+ monotonic_range_key_projection(sort_expr, mapping,
input_eq_properties)
+ .find(|(_, source)| {
+ monotonic_fn_keeps_partitions_disjoint(
+ source,
+ &sort_expr.expr,
+ &split_points,
+ key_idx,
+ sort_expr.options.descending,
+ )
+ })?;
+ // Fail closed: if the split cannot be rewritten into the
+ // transformed domain, do not keep Range with source-domain bounds.
+ split_points = project_split_points_through_fn(
+ &source,
+ &sort_expr.expr,
+ &split_points,
+ key_idx,
+ )?;
+ 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,
+ })
+ }
+}
+
+/// Yields projection mappings whose source is a same-direction monotonic
+/// transform of `sort_expr`. Callers pick the first candidate that also
+/// keeps adjacent partitions disjoint.
+fn monotonic_range_key_projection<'a>(
+ sort_expr: &'a PhysicalSortExpr,
+ mapping: &'a ProjectionMapping,
+ eq_properties: &'a EquivalenceProperties,
+) -> impl Iterator<Item = (Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)> + 'a
{
+ mapping
+ .iter()
+ .filter(|(source, _)| {
+ eq_properties.check_monotonic_transform(source, &sort_expr.expr)
+ })
+ .map(|(source, targets)| (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 on the open side
+/// of the split: predecessor for ASC, successor for DESC.
+fn monotonic_fn_keeps_partitions_disjoint(
+ fn_expr: &Arc<dyn PhysicalExpr>,
+ range_key: &Arc<dyn PhysicalExpr>,
+ split_points: &[SplitPoint],
+ key_idx: usize,
+ descending: bool,
+) -> bool {
+ split_points.iter().all(|split_point| {
+ let Some(split_value) = split_point.values().get(key_idx) else {
+ return false;
+ };
+ // ASC: partition 0 is key < split, adjacent is predecessor.
+ // DESC: partition 0 is key > split, adjacent is successor.
+ let adjacent = if descending {
+ checked_successor(split_value)
+ } else {
+ checked_predecessor(split_value)
+ };
+ let Some(adjacent) = adjacent else {
+ return false;
+ };
+ let Some(at_split) = evaluate_expr_on_key(fn_expr, range_key,
split_value) else {
+ return false;
+ };
+ let Some(across_split) = evaluate_expr_on_key(fn_expr, range_key,
&adjacent)
+ else {
+ return false;
+ };
+ at_split != across_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> {
Review Comment:
We might need type check at the beginning of the evaluation
```rust
// The substituted literal must carry exactly the range key's type: a
// different time unit or timezone silently changes what date_bin /
// date_trunc compute, and would make the disjointness check answer a
// question about a function the data was never partitioned by.
if range_key.data_type(schema).ok()? != value.data_type() {
return None;
}
```
--
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]
