namanjain24-sudo commented on code in PR #25781:
URL: https://github.com/apache/datafusion/pull/25781#discussion_r4226720526
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -828,6 +912,45 @@ fn can_pullup_over_aggregation(expr: &Expr) -> bool {
}
}
+/// Whether `expr` computes `GROUPING`/`GROUPING_ID`, which reads whether a
+/// row's grouping set groups by a particular column.
+fn is_grouping_call(expr: &Expr) -> bool {
+ matches!(expr, Expr::AggregateFunction(agg) if
agg.func.name().eq_ignore_ascii_case("grouping"))
+}
+
+/// Columns read by a node strictly above the first `Aggregate` found in
+/// `plan`, which is the subquery's original, not yet rewritten, plan.
+///
+/// [`PullUpCorrelatedExpr::f_up`] runs bottom-up, so by the time it visits an
+/// `Aggregate` it cannot yet tell whether a `HAVING` or `Projection` above it
+/// reads one of the columns a grouping set pull up would add. This walks the
+/// plan top-down instead, before the rewrite starts, and stops at the first
+/// `Aggregate` along each branch, collecting the columns every node above it
+/// reads in its own expressions. A nested `Subquery` is a different
+/// correlation scope and is skipped, the same way
+/// [`PullUpCorrelatedExpr::f_down`] skips it.
+fn columns_read_above_aggregate(plan: &LogicalPlan) -> BTreeSet<Column> {
+ fn walk(plan: &LogicalPlan, above: &mut BTreeSet<Column>) -> bool {
+ if matches!(plan, LogicalPlan::Subquery(_)) {
+ return false;
+ }
+ if matches!(plan, LogicalPlan::Aggregate(_)) {
Review Comment:
test
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -828,6 +912,45 @@ fn can_pullup_over_aggregation(expr: &Expr) -> bool {
}
}
+/// Whether `expr` computes `GROUPING`/`GROUPING_ID`, which reads whether a
+/// row's grouping set groups by a particular column.
+fn is_grouping_call(expr: &Expr) -> bool {
+ matches!(expr, Expr::AggregateFunction(agg) if
agg.func.name().eq_ignore_ascii_case("grouping"))
+}
+
+/// Columns read by a node strictly above the first `Aggregate` found in
+/// `plan`, which is the subquery's original, not yet rewritten, plan.
+///
+/// [`PullUpCorrelatedExpr::f_up`] runs bottom-up, so by the time it visits an
+/// `Aggregate` it cannot yet tell whether a `HAVING` or `Projection` above it
+/// reads one of the columns a grouping set pull up would add. This walks the
+/// plan top-down instead, before the rewrite starts, and stops at the first
+/// `Aggregate` along each branch, collecting the columns every node above it
+/// reads in its own expressions. A nested `Subquery` is a different
+/// correlation scope and is skipped, the same way
+/// [`PullUpCorrelatedExpr::f_down`] skips it.
+fn columns_read_above_aggregate(plan: &LogicalPlan) -> BTreeSet<Column> {
+ fn walk(plan: &LogicalPlan, above: &mut BTreeSet<Column>) -> bool {
+ if matches!(plan, LogicalPlan::Subquery(_)) {
+ return false;
+ }
+ if matches!(plan, LogicalPlan::Aggregate(_)) {
Review Comment:
Applied exactly as suggested. Verified by reverting locally and running the
three repros: query 1 returns count=0 for every row, queries 2 and 3 return
false for every row — exactly the wrong output described. With the fix, all
three correctly error instead, and the existing ((k), (j)) EXISTS case still
decorrelates. Added the three as statement error cases in subquery.slt; full
suite passes.
--
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]
