jayzhan211 commented on code in PR #25781:
URL: https://github.com/apache/datafusion/pull/25781#discussion_r4294887433
##########
datafusion/optimizer/src/decorrelate.rs:
##########
@@ -136,6 +144,36 @@ impl PullUpCorrelatedExpr {
self.exists_sub_query = exists_sub_query;
self
}
+
+ /// Record the columns read above the subquery's `Aggregate`, computed
+ /// from `subquery_plan` (the plan as it is before this rewrite runs).
+ ///
+ /// A grouping set that omits a correlated column only needs to reject the
+ /// pull up when something reads the column the pull up would fill in, so
+ /// [`Self::f_up`] needs this computed ahead of time. See
+ /// `columns_read_above_aggregate`.
+ pub fn with_column_refs_above_aggregate(
Review Comment:
`PullUpCorrelatedExpr` is pub, so a caller that uses `new()` without
`with_column_refs_above_aggregate` gets an empty set → `SafeToExtend` for
`HAVING k IS NULL` / `GROUPING(j)` (rejected before this PR). A stack filled in
`f_down` gives `f_up` the ancestors' reads directly: no pre-pass, no builder,
no call-site changes (+23/−89; full slt, optimizer tests, clippy and all
earlier repros green):
```diff
- column_refs_above_aggregate: BTreeSet<Column>,
+ /// Columns each node on the path from the root reads in its own
+ /// expressions; pushed in `f_down`, popped in `f_up`.
+ column_refs_above: Vec<BTreeSet<Column>>,
```
```diff
fn f_down(&mut self, plan: LogicalPlan) ->
Result<Transformed<LogicalPlan>> {
+ let mut column_refs = BTreeSet::new();
+ for expr in plan.expressions() {
+ column_refs.extend(expr.column_refs().into_iter().cloned());
+ }
+ self.column_refs_above.push(column_refs);
```
```diff
fn f_up(&mut self, plan: LogicalPlan) ->
Result<Transformed<LogicalPlan>> {
+ self.column_refs_above.pop();
```
```diff
- let reads_grouping_id = self
- .column_refs_above_aggregate
- .iter()
- .any(|col| col.name == Aggregate::INTERNAL_GROUPING_ID);
- let read_above = required_cols
- .iter()
- .any(|col| self.column_refs_above_aggregate.contains(*col));
-
- if has_empty_set || reads_grouping_id || read_above {
+ let unsafe_read = self.column_refs_above.iter().flatten().any(|col|
{
+ col.name == Aggregate::INTERNAL_GROUPING_ID ||
required_cols.contains(col)
+ });
+
+ if has_empty_set || unsafe_read {
```
Then delete `with_column_refs_above_aggregate`,
`columns_read_above_aggregate`, and the `.with_column_refs_above_aggregate(..)`
calls in `scalar_subquery_to_join.rs`, `decorrelate_predicate_subquery.rs` and
`decorrelate_lateral_join.rs`.
--
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]
