NGA-TRAN commented on code in PR #24501:
URL: https://github.com/apache/datafusion/pull/24501#discussion_r3849090348
##########
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:
Done
##########
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:
Fixed. We now require SortProperties::Ordered(options) to match the source
SortOptions, so -x is false. Covered by
`check_monotonic_transform_rejects_order_reversing_negation`.
##########
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:
Done
##########
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:
Good catch — allow_subset is Hash-key subsetting, not a monotonic transform
of the range key. `range_monotonic_fn_satisfies_keys` now runs even when
allow_subset is false and still returns Subset
##########
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:
Added in `test_range_partitioning_project_keeps_range_key`:
SELECT a, date_bin(a) → Range([a]) (the range key is kept; date_bin is extra)
SELECT a AS b, date_bin(a) AS bucket → Range([b]) (alias of a)
Both stay on the original split. The monotonic-fn path only runs when the
projection drops a.
##########
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:
Tracked in #24644
##########
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:
Done
##########
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:
Agreed on first target — it is just the alias.
`monotonic_range_key_projection` now yields all same-direction monotonic
sources. RangePartitioning::project takes the first that also passes
`monotonic_fn_keeps_partitions_disjoint`, so a straddling first bin no longer
forces `Unknown` if a later bin is aligned.
Covered in
`test_range_partitioning_project_skips_non_monotonic_and_straddling_bins`:
SELECT -a, date_bin(60s, a) → skips -a, keeps Range on the 60s bin
SELECT -a, date_bin AS bin1, date_bin AS bin2 → uses bin1
SELECT date_bin(70s, a), date_bin(60s, a) → 70s straddles this hour split,
60s does not, so Range is preserved through the 60s bin
Used 70s instead of 45s because 45s is aligned at 2024-01-01T01:00:00 (it
divides the split), so it would not demonstrate the “first straddles, second
does not” case.
--
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]
