[arrow-datafusion] branch main updated: minor: add with_plan for Subquery (#5680)

[jakevin]

This is an automated email from the ASF dual-hosted git repository.

jakevin pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 86740a8048 minor: add with_plan for Subquery (#5680)
86740a8048 is described below

commit 86740a8048c4766f1a15944cb621013aceb33cfc
Author: jakevin <jakevin...@gmail.com>
AuthorDate: Wed Mar 22 17:30:18 2023 +0800

    minor: add with_plan for Subquery (#5680)
---
 datafusion/expr/src/logical_plan/plan.rs             |  7 +++++++
 datafusion/optimizer/src/decorrelate_where_exists.rs |  8 ++------
 datafusion/optimizer/src/decorrelate_where_in.rs     | 13 ++++++-------
 datafusion/optimizer/src/scalar_subquery_to_join.rs  |  7 ++-----
 4 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/datafusion/expr/src/logical_plan/plan.rs 
b/datafusion/expr/src/logical_plan/plan.rs
index da490db395..58faf69487 100644
--- a/datafusion/expr/src/logical_plan/plan.rs
+++ b/datafusion/expr/src/logical_plan/plan.rs
@@ -1946,6 +1946,13 @@ impl Subquery {
             _ => plan_err!("Could not coerce into ScalarSubquery!"),
         }
     }
+
+    pub fn with_plan(&self, plan: Arc<LogicalPlan>) -> Subquery {
+        Subquery {
+            subquery: plan,
+            outer_ref_columns: self.outer_ref_columns.clone(),
+        }
+    }
 }
 
 impl Debug for Subquery {
diff --git a/datafusion/optimizer/src/decorrelate_where_exists.rs 
b/datafusion/optimizer/src/decorrelate_where_exists.rs
index 35d7858ba3..88117c6be9 100644
--- a/datafusion/optimizer/src/decorrelate_where_exists.rs
+++ b/datafusion/optimizer/src/decorrelate_where_exists.rs
@@ -62,12 +62,8 @@ impl DecorrelateWhereExists {
                         .try_optimize(&subquery.subquery, config)?
                         .map(Arc::new)
                         .unwrap_or_else(|| subquery.subquery.clone());
-                    let subquery = Subquery {
-                        subquery: subquery_plan,
-                        outer_ref_columns: subquery.outer_ref_columns.clone(),
-                    };
-                    let subquery = SubqueryInfo::new(subquery.clone(), 
*negated);
-                    subqueries.push(subquery);
+                    let new_subquery = subquery.with_plan(subquery_plan);
+                    subqueries.push(SubqueryInfo::new(new_subquery, *negated));
                 }
                 _ => others.push((*it).clone()),
             }
diff --git a/datafusion/optimizer/src/decorrelate_where_in.rs 
b/datafusion/optimizer/src/decorrelate_where_in.rs
index 8d7988d711..d89fd935be 100644
--- a/datafusion/optimizer/src/decorrelate_where_in.rs
+++ b/datafusion/optimizer/src/decorrelate_where_in.rs
@@ -69,13 +69,12 @@ impl DecorrelateWhereIn {
                         .try_optimize(&subquery.subquery, config)?
                         .map(Arc::new)
                         .unwrap_or_else(|| subquery.subquery.clone());
-                    let subquery = Subquery {
-                        subquery: subquery_plan,
-                        outer_ref_columns: subquery.outer_ref_columns.clone(),
-                    };
-                    let subquery =
-                        SubqueryInfo::new(subquery.clone(), (**expr).clone(), 
*negated);
-                    subqueries.push(subquery);
+                    let new_subquery = subquery.with_plan(subquery_plan);
+                    subqueries.push(SubqueryInfo::new(
+                        new_subquery,
+                        (**expr).clone(),
+                        *negated,
+                    ));
                     // TODO: if subquery doesn't get optimized, optimized 
children are lost
                 }
                 _ => others.push((*it).clone()),
diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs 
b/datafusion/optimizer/src/scalar_subquery_to_join.rs
index 303d97bb4f..df0b9245fa 100644
--- a/datafusion/optimizer/src/scalar_subquery_to_join.rs
+++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs
@@ -76,11 +76,8 @@ impl ScalarSubqueryToJoin {
                                 .try_optimize(&subquery.subquery, config)?
                                 .map(Arc::new)
                                 .unwrap_or_else(|| subquery.subquery.clone());
-                            let subquery = Subquery {
-                                subquery: subquery_plan,
-                                outer_ref_columns: 
subquery.outer_ref_columns.clone(),
-                            };
-                            let res = SubqueryInfo::new(subquery, expr, *op, 
lhs);
+                            let new_subquery = 
subquery.with_plan(subquery_plan);
+                            let res = SubqueryInfo::new(new_subquery, expr, 
*op, lhs);
                             subqueries.push(res);
                             Ok(())
                         };