comphead commented on code in PR #26013:
URL: https://github.com/apache/datafusion/pull/26013#discussion_r4296832029
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -783,6 +798,51 @@ fn correlated_filter_columns(plan: &LogicalPlan) ->
Vec<Column> {
cols
}
+/// Adds `cols` to the `PARTITION BY` of every window function of `window`. A
+/// Projection on top gives the window columns their old names back.
+fn partition_window_by(window: &Window, cols: &BTreeSet<Column>) ->
Result<LogicalPlan> {
+ let new_window_expr = window
+ .window_expr
+ .iter()
+ .map(|expr| add_partition_by(expr, cols))
+ .collect::<Vec<_>>();
+ let mut proj_exprs = window
+ .input
+ .schema()
+ .columns()
+ .into_iter()
+ .map(Expr::Column)
+ .collect::<Vec<_>>();
+ for (old, new) in window.window_expr.iter().zip(&new_window_expr) {
+ let new_col =
Expr::Column(Column::from_name(new.schema_name().to_string()));
Review Comment:
Fixed in 0092bee2a. Each extended window expression is now aliased to its
old name, so the new `Window` has no duplicate field. A test covers a lower
window that already outputs the extended name.
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -783,6 +798,51 @@ fn correlated_filter_columns(plan: &LogicalPlan) ->
Vec<Column> {
cols
}
+/// Adds `cols` to the `PARTITION BY` of every window function of `window`. A
+/// Projection on top gives the window columns their old names back.
+fn partition_window_by(window: &Window, cols: &BTreeSet<Column>) ->
Result<LogicalPlan> {
+ let new_window_expr = window
+ .window_expr
+ .iter()
+ .map(|expr| add_partition_by(expr, cols))
+ .collect::<Vec<_>>();
+ let mut proj_exprs = window
+ .input
+ .schema()
+ .columns()
+ .into_iter()
+ .map(Expr::Column)
+ .collect::<Vec<_>>();
+ for (old, new) in window.window_expr.iter().zip(&new_window_expr) {
+ let new_col =
Expr::Column(Column::from_name(new.schema_name().to_string()));
+
proj_exprs.push(new_col.alias_if_changed(old.schema_name().to_string())?);
+ }
+ LogicalPlanBuilder::from(Arc::clone(&window.input))
+ .window(new_window_expr)?
+ .project(proj_exprs)?
Review Comment:
Fixed in 0092bee2a, same change as above. The extended expression keeps its
old output name.
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -783,6 +798,51 @@ fn correlated_filter_columns(plan: &LogicalPlan) ->
Vec<Column> {
cols
}
+/// Adds `cols` to the `PARTITION BY` of every window function of `window`. A
+/// Projection on top gives the window columns their old names back.
+fn partition_window_by(window: &Window, cols: &BTreeSet<Column>) ->
Result<LogicalPlan> {
+ let new_window_expr = window
+ .window_expr
+ .iter()
+ .map(|expr| add_partition_by(expr, cols))
+ .collect::<Vec<_>>();
+ let mut proj_exprs = window
Review Comment:
Done in 0092bee2a. `partition_window_by` now builds `Window::try_new` with
`alias_if_changed(old.schema_name())`, and the extra `Projection` is gone.
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -783,6 +798,51 @@ fn correlated_filter_columns(plan: &LogicalPlan) ->
Vec<Column> {
cols
}
+/// Adds `cols` to the `PARTITION BY` of every window function of `window`. A
+/// Projection on top gives the window columns their old names back.
+fn partition_window_by(window: &Window, cols: &BTreeSet<Column>) ->
Result<LogicalPlan> {
+ let new_window_expr = window
+ .window_expr
+ .iter()
+ .map(|expr| add_partition_by(expr, cols))
+ .collect::<Vec<_>>();
+ let mut proj_exprs = window
+ .input
+ .schema()
+ .columns()
+ .into_iter()
+ .map(Expr::Column)
+ .collect::<Vec<_>>();
+ for (old, new) in window.window_expr.iter().zip(&new_window_expr) {
+ let new_col =
Expr::Column(Column::from_name(new.schema_name().to_string()));
Review Comment:
Fixed in 0092bee2a. An `Alias` keeps its `relation` because
`add_partition_by` clones it, and the restoring `Projection` that dropped it no
longer exists.
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -577,6 +581,17 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
.iter()
.all(|expr| partitions_by_all(expr,
&local_correlated_cols))
{
+ // `inner.y = outer.k` keeps the rows of one `y` value for
+ // each outer row, so partitioning by `y` as well gives
each
+ // window function the same rows.
+ if self.can_extend_partition_by
+ &&
!self.collected_count_expr_map.contains_key(&*window.input)
Review Comment:
I dropped the explicit guard in 0092bee2a. `pass_through_pulled_up_cols`
already sets `can_pull_up = false` when a node input is in
`collected_count_expr_map`. `Projection`, `SubqueryAlias`, `Limit` and `Filter`
propagate that entry upward. A nested `Window` stops at the inner one. The
existing slt case for a window over a count bug aggregate still stays
correlated.
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -828,6 +888,24 @@ fn can_pullup_over_aggregation(expr: &Expr) -> bool {
}
}
+/// Whether `expr` is `col = <expr>` or `<expr> = col`, where `<expr>` reads no
+/// subquery column. `CAST(col AS ..)` does not qualify: a lossy cast lets
+/// several `col` values match one outer row, see #26014.
+fn is_column_equality(expr: &Expr) -> bool {
+ let Expr::BinaryExpr(BinaryExpr {
+ left,
+ op: Operator::Eq,
+ right,
+ }) = expr
+ else {
+ return false;
+ };
+ match (&**left, &**right) {
+ (Expr::Column(_), other) | (other, Expr::Column(_)) =>
!other.any_column_refs(),
Review Comment:
Fixed in 0092bee2a. `is_partitionable_equality` rejects a volatile outer
expression, and a new slt case covers it. It fails if the check is removed.
--
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]
