gene-bordegaray commented on code in PR #24501:
URL: https://github.com/apache/datafusion/pull/24501#discussion_r3940437705
##########
datafusion/sqllogictest/test_files/range_sorted_time_bin_agg.slt:
##########
@@ -146,6 +156,172 @@ k2 2024-01-01T00:30:00 7
k2 2024-01-01T01:30:00 30
k2 2024-01-01T01:45:00 5
+##########
Review Comment:
thanks for tehse diagrams, these are awesome
##########
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> {
+ 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).
+///
+/// - [`PartitioningSatisfaction::Exact`] when every required key is such an
+/// `f(x)` (for example `GROUP BY date_bin(x)`).
+/// - [`PartitioningSatisfaction::Subset`] when at least one required key is
+/// such an `f(x)` and others are not (for example `GROUP BY date_bin(x),
baz`),
+/// and `allow_subset` is true.
+///
+/// Multi-key Range + monotonic transforms (e.g. `Range([key, timestamp])`
+/// + `GROUP BY (key, date_bin(timestamp))`) is not handled here; see #24644.
+fn range_monotonic_fn_satisfaction(
+ range: &RangePartitioning,
+ required_exprs: &[Arc<dyn PhysicalExpr>],
+ eq_properties: &EquivalenceProperties,
+ allow_subset: bool,
+) -> PartitioningSatisfaction {
+ if range.ordering().len() != 1 || required_exprs.is_empty() {
+ return PartitioningSatisfaction::NotSatisfied;
+ }
+ let range_sort = &range.ordering()[0];
+ let range_key = &range_sort.expr;
+ let matching = required_exprs
+ .iter()
+ .filter(|required| {
+ eq_properties.check_monotonic_transform(required, range_key)
+ && monotonic_fn_keeps_partitions_disjoint(
+ required,
+ range_key,
+ range.split_points(),
+ 0,
+ range_sort.options.descending,
+ )
+ })
+ .count();
+
+ if matching == 0 {
+ PartitioningSatisfaction::NotSatisfied
+ } else if matching == required_exprs.len() {
Review Comment:
this may not be safe for co-partitioning joins. is this being used for this?
If so I think it should be diabled in this PR
##########
datafusion/physical-expr/src/equivalence/properties/mod.rs:
##########
@@ -1319,6 +1320,33 @@ impl EquivalenceProperties {
.unwrap_or_else(|_| ExprProperties::new_unknown())
}
+ /// Returns true when `transformed_expr` is a same-direction monotonic
+ /// transform of `source_expr` plus literals, such as
+ /// `date_bin(interval, timestamp)` or `date_trunc(unit, timestamp)`.
+ ///
+ /// The identity `transformed_expr == source_expr` returns false so callers
+ /// can treat "emit the key as-is" separately from "emit a function of the
+ /// key". Order-reversing transforms such as `-x` also return false.
+ pub(crate) fn check_monotonic_transform(
+ &self,
+ transformed_expr: &Arc<dyn PhysicalExpr>,
+ source_expr: &Arc<dyn PhysicalExpr>,
+ ) -> bool {
+ if transformed_expr.eq(source_expr) {
+ return false;
+ }
+ let options = SortOptions::default();
Review Comment:
I dont belive sort options should be deafult in all cases here
##########
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> {
+ 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).
+///
+/// - [`PartitioningSatisfaction::Exact`] when every required key is such an
+/// `f(x)` (for example `GROUP BY date_bin(x)`).
+/// - [`PartitioningSatisfaction::Subset`] when at least one required key is
+/// such an `f(x)` and others are not (for example `GROUP BY date_bin(x),
baz`),
+/// and `allow_subset` is true.
+///
+/// Multi-key Range + monotonic transforms (e.g. `Range([key, timestamp])`
+/// + `GROUP BY (key, date_bin(timestamp))`) is not handled here; see #24644.
+fn range_monotonic_fn_satisfaction(
+ range: &RangePartitioning,
+ required_exprs: &[Arc<dyn PhysicalExpr>],
+ eq_properties: &EquivalenceProperties,
+ allow_subset: bool,
+) -> PartitioningSatisfaction {
+ if range.ordering().len() != 1 || required_exprs.is_empty() {
+ return PartitioningSatisfaction::NotSatisfied;
+ }
+ let range_sort = &range.ordering()[0];
+ let range_key = &range_sort.expr;
+ let matching = required_exprs
+ .iter()
+ .filter(|required| {
+ eq_properties.check_monotonic_transform(required, range_key)
+ && monotonic_fn_keeps_partitions_disjoint(
+ required,
+ range_key,
+ range.split_points(),
+ 0,
+ range_sort.options.descending,
+ )
+ })
+ .count();
+
+ if matching == 0 {
+ PartitioningSatisfaction::NotSatisfied
+ } else if matching == required_exprs.len() {
Review Comment:
```text
left: Range(x), split 10
right: Range(y), split 10
join: x + 1 = y
```
The left input independently satisfies KeyPartitioned(x + 1), but its join
is 11. Then the right boundary is 10. `compatible_co_partitioning_layout` still
compares the original split tuples (10 == 10), so the partitioned join is
accepted. Then left.x = 9 and right.y = 10 have the same join key but occur in
partitions 0 and 1 and never meet.
--
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]
