Add OuterReferenceColumn to Expr to represent correlated expression #5593

mingmwang · 2023-03-14T11:28:37Z

Which issue does this PR close?

Closes #5571.

Original design doc:
https://docs.google.com/document/d/1j5vHyva-T_5l3POnHSS0-r5TOp86sEkfSUcRLqJZaK4/edit#heading=h.wgtywgh0byuf

Rationale for this change

This is first part to support more general and complex subqueries and subquery expressions.
This PR add the OuterReferenceColumn to Expr to represent the correlated expressions, no matter how deeply
the correlated expressions exist.

During the SqlToRel(planner) process, correlated expressions will be marked as OuterReferenceColumn.
This will make sure the correlated expressions can pass the inner plan's schema validation. It will also simply the logic
of Decorrelation rules. The goal of the Decorrelation rule is just to find out the correlated expressions and convert to different Joins.

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

mingmwang · 2023-03-14T11:50:25Z

I will add more UTs to cover more complex cases:

Subquery with Limit
Subquery with Order by
Subquery with Window Functions
Subquery with Union

mingmwang · 2023-03-16T09:11:45Z

@jackwener @ygf11 @alamb

ygf11

Looks great to me. Thanks @mingmwang.

ygf11 · 2023-03-17T08:45:51Z

datafusion/expr/src/logical_plan/plan.rs

@@ -1919,14 +1947,11 @@ impl Join {
 pub struct Subquery {
    /// The subquery
    pub subquery: Arc<LogicalPlan>,
+    /// The outer references used in the subquery
+    pub outer_ref_columns: Vec<Expr>,


👍 This field can help to check if it is a correlated subquery.

ygf11 · 2023-03-17T08:50:48Z

datafusion/sql/src/expr/subquery.rs

+        planner_context.outer_query_schema = Some(input_schema.clone());
+        let sub_plan = self.query_to_plan(subquery, planner_context)?;
+        let outer_ref_columns = sub_plan.all_out_ref_exprs();
+        planner_context.outer_query_schema = old_outer_query_schema;


Can we support the following subquery in the future?

# three level # the innermost(`t3`) reference outermost(`t1`) column SELECT * FROM t1 WHERE EXISTS ( SELECT * FROM t2 WHERE t1_id = t2_id AND EXISTS ( SELECT * FROM t3 WHERE t3_id > t1_id ) );

I just think the optimizer rule you plan to add(#5492) can also optimize it.

I'm afraid I will not support this in recent. I think the method mentioned in the "Unnesting Arbitrary Queries" paper can cover this, but there is additional complexity need to deal with if there are some ambiguous to the out ref columns, because different level of out plans' schema might have ambiguous columns.

alamb

Looks good to me @mingmwang -- I left some suggestions about cleaning up the code in the planner, but the overall idea looks great to me. Thank you

alamb · 2023-03-17T10:53:02Z

datafusion/core/src/datasource/listing/helpers.rs

@@ -81,6 +81,7 @@ impl ExpressionVisitor for ApplicabilityVisitor<'_> {
            }
            Expr::Literal(_)
            | Expr::Alias(_, _)
+            | Expr::OuterReferenceColumn(_, _)


This idea makes a lot of sense to me 👍

alamb · 2023-03-17T10:54:01Z

datafusion/core/tests/sql/subqueries.rs

@@ -179,3 +179,117 @@ async fn in_subquery_with_same_table() -> Result<()> {

    Ok(())
 }
+
+#[tokio::test]
+async fn support_agg_correlated_columns() -> Result<()> {


if possible I would recommend adding these tests as sqllogictest -- https://github.com/apache/arrow-datafusion/tree/main/datafusion/core/tests/sqllogictests

They are much easier to change / maintain in my opinion

If test just relate with plan tree, I think it's also ok to keep it here.
Because we can debug it easily, and see plan directly in the code.

alamb · 2023-03-17T10:54:32Z

datafusion/core/tests/sql/subqueries.rs

+#[tokio::test]
+async fn support_join_correlated_columns2() -> Result<()> {
+    let ctx = create_sub_query_join_context("t0_id", "t1_id", "t2_id", true)?;
+    let sql = "SELECT t0_id, t0_name FROM t0 WHERE EXISTS (SELECT 1 FROM t1 INNER JOIN (select * from t2 where t2.t2_name = t0.t0_name) as t2 ON(t1.t1_id = t2.t2_id ))";


This is neat to see a correlation to the (second) layer of query

alamb · 2023-03-17T10:57:21Z

datafusion/optimizer/src/utils.rs

+        if filter.contains_outer() {
+            joins.push(strip_outer_reference((*filter).clone()));
+        } else {
+            // TODO remove the logic


I don't understand this TODO -- is it meant for this PR or for some other one

In the original decorrelation rules, there was complex logic to find out join columns, this is unnecessary after this PR.
I will clean up the code in the next PR.

alamb · 2023-03-17T11:01:03Z

datafusion/sql/src/expr/identifier.rs

+                    let outer_query_schema_opt =
+                        planner_context.outer_query_schema.clone();
+                    if let Some(outer) = outer_query_schema_opt.as_ref() {


Is there a reason for this clone (of the entire schema)? Could it be more like

Suggested change

let outer_query_schema_opt =

planner_context.outer_query_schema.clone();

if let Some(outer) = outer_query_schema_opt.as_ref() {

let outer_query_schema_opt = planner_context.outer_query_schema.as_ref();

if let Some(outer) = outer_query_schema_opt {

The same thing applies to the other places in this file

alamb · 2023-03-17T11:02:08Z

datafusion/sql/src/expr/identifier.rs

+                        let outer_query_schema_opt =
+                            planner_context.outer_query_schema.clone();
+                        if let Some(outer) = outer_query_schema_opt.as_ref() {
+                            match outer.field_with_qualified_name(&relation, &name) {
+                                Ok(field) => {
+                                    // found an exact match on a qualified name in the outer plan schema, so this is an outer reference column
+                                    Ok(Expr::OuterReferenceColumn(
+                                        field.data_type().clone(),
+                                        Column {
+                                            relation: Some(relation),
+                                            name,
+                                        },
+                                    ))
+                                }
+                                Err(_) => Ok(Expr::Column(Column {
+                                    relation: Some(relation),
+                                    name,
+                                })),
+                            }
+                        } else {
+                            // table.column identifier
+                            Ok(Expr::Column(Column {
+                                relation: Some(relation),
+                                name,
+                            }))
+                        }


This pattern appears to be repeated several times -- maybe it could be refactored into a reusable function rather than being inlined several places?

Sure, will do.

alamb · 2023-03-17T11:03:14Z

datafusion/sql/src/select.rs

-                } else {
-                    vec![]
-                };
+                let outer_query_schema = planner_context.outer_query_schema.clone();


again here I think we can avoid a clone

Agree with it.
We can avoid a clone

alamb · 2023-03-17T11:03:30Z

datafusion/sql/tests/integration_test.rs

@@ -2635,7 +2635,7 @@ fn exists_subquery() {
        \n  Filter: EXISTS (<subquery>)\
        \n    Subquery:\
        \n      Projection: person.first_name\
-        \n        Filter: person.last_name = p.last_name AND person.state = p.state\
+        \n        Filter: person.last_name = outer_ref(p.last_name) AND person.state = outer_ref(p.state)\


I find this actually much easier to read 👍

jackwener

A nice job to me.

jackwener · 2023-03-17T13:40:54Z

datafusion/sql/src/select.rs

-                } else {
-                    vec![]
-                };
+                let outer_query_schema = planner_context.outer_query_schema.clone();


Agree with it.
We can avoid a clone

jackwener · 2023-03-17T15:13:51Z

datafusion/optimizer/src/decorrelate_where_in.rs

@@ -149,6 +152,7 @@ fn optimize_where_in(
    let projection = Projection::try_from_plan(&query_info.query.subquery)
        .map_err(|e| context!("a projection is required", e))?;
    let subquery_input = projection.input.clone();
+    // TODO add the validate logic to Analyzer


😍Look forward it!

In the future, I also prepare to do some check job in the Analyzer.

jackwener · 2023-03-17T15:17:37Z

datafusion/core/tests/sql/subqueries.rs

@@ -179,3 +179,117 @@ async fn in_subquery_with_same_table() -> Result<()> {

    Ok(())
 }
+
+#[tokio::test]
+async fn support_agg_correlated_columns() -> Result<()> {


If test just relate with plan tree, I think it's also ok to keep it here.
Because we can debug it easily, and see plan directly in the code.

alamb · 2023-03-18T11:38:56Z

Looks like this PR has some conflicts to fix but then it will be good to go

…hema

…d columns in subquery

mingmwang · 2023-03-20T09:14:37Z

@alamb @jackwener
Please help to take a look again. Conflicts are resolved now.

alamb · 2023-03-20T13:26:21Z

Please help to take a look again. Conflicts are resolved now.

LGTM but sadly there appear to be some more conflicts now:

I took the liberty to merge from main and resolve the new conflicts -- I plan to merge this PR when CI passes

Add OuterReferenceColumn

a2d6cae

github-actions bot added core Core DataFusion crate logical-expr Logical plan and expressions optimizer Optimizer rules sql SQL Planner labels Mar 14, 2023

fix intg test

0cc8bda

jackwener self-requested a review March 16, 2023 10:43

ygf11 approved these changes Mar 17, 2023

View reviewed changes

alamb approved these changes Mar 17, 2023

View reviewed changes

jackwener approved these changes Mar 17, 2023

View reviewed changes

mingmwang added 4 commits March 20, 2023 13:47

resolve conflicts

788d3c4

resolve review comments, avoid unncessary clone to the outer_query_sc…

e9d2e94

…hema

add more UT to cover more complex subqueryies, fix order by correlate…

dfc65b8

…d columns in subquery

fix cargo fmt

4b0b494

alamb added 2 commits March 20, 2023 09:23

Merge remote-tracking branch 'apache/main' into subquery_out_column

bf15ec3

fix: logical merge conflict

fc8d537

alamb merged commit 26e1b20 into apache:main Mar 20, 2023

jonahgao mentioned this pull request Jul 23, 2024

Parsing SQL strings to Exprs with the qualified schema #11562

Merged

Add OuterReferenceColumn to Expr to represent correlated expression #5593

Add OuterReferenceColumn to Expr to represent correlated expression #5593

Uh oh!

Conversation

mingmwang commented Mar 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

mingmwang commented Mar 14, 2023

Uh oh!

mingmwang commented Mar 16, 2023

Uh oh!

ygf11 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jackwener left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jackwener Mar 17, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb commented Mar 18, 2023

Uh oh!

mingmwang commented Mar 20, 2023

Uh oh!

alamb commented Mar 20, 2023

Uh oh!

Uh oh!

mingmwang commented Mar 14, 2023 •

edited

Loading

jackwener Mar 17, 2023 •

edited

Loading