refactor: move collect_subquery_cols() to common utils crate

Author: wiedld

datafusion/expr/src/logical_plan/invariants.rs

@@ -15,20 +15,18 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::collections::BTreeSet;
-
 use recursive::recursive;
 
 use datafusion_common::{
     plan_err,
     tree_node::{TreeNode, TreeNodeRecursion},
-    Column, DFSchema, DFSchemaRef, DataFusionError, Result,
+    DFSchemaRef, DataFusionError, Result,
 };
 
 use crate::{
     expr::{Exists, InSubquery},
     expr_rewriter::strip_outer_reference,
-    utils::split_conjunction,
+    utils::{collect_subquery_cols, split_conjunction},
     Aggregate, Expr, Filter, Join, JoinType, LogicalPlan, Window,
 };
 
@@ -383,23 +381,6 @@ fn check_mixed_out_refer_in_window(window: &Window) -> Result<()> {
     }
 }
 
-fn collect_subquery_cols(
-    exprs: &[Expr],
-    subquery_schema: &DFSchema,
-) -> Result<BTreeSet<Column>> {
-    exprs.iter().try_fold(BTreeSet::new(), |mut cols, expr| {
-        let mut using_cols: Vec<Column> = vec![];
-        for col in expr.column_refs().into_iter() {
-            if subquery_schema.has_column(col) {
-                using_cols.push(col.clone());
-            }
-        }
-
-        cols.extend(using_cols);
-        Result::<_>::Ok(cols)
-    })
-}
-
 #[cfg(test)]
 mod test {
     use std::cmp::Ordering;

datafusion/expr/src/utils.rs

@@ -18,7 +18,7 @@
 //! Expression utilities
 
 use std::cmp::Ordering;
-use std::collections::HashSet;
+use std::collections::{BTreeSet, HashSet};
 use std::ops::Deref;
 use std::sync::Arc;
 
@@ -1402,6 +1402,24 @@ pub fn format_state_name(name: &str, state_name: &str) -> String {
     format!("{name}[{state_name}]")
 }
 
+/// Determine the set of [`Column`]s produced by the subquery.
+pub fn collect_subquery_cols(
+    exprs: &[Expr],
+    subquery_schema: &DFSchema,
+) -> Result<BTreeSet<Column>> {
+    exprs.iter().try_fold(BTreeSet::new(), |mut cols, expr| {
+        let mut using_cols: Vec<Column> = vec![];
+        for col in expr.column_refs().into_iter() {
+            if subquery_schema.has_column(col) {
+                using_cols.push(col.clone());
+            }
+        }
+
+        cols.extend(using_cols);
+        Result::<_>::Ok(cols)
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

datafusion/optimizer/src/decorrelate.rs

@@ -22,15 +22,16 @@ use std::ops::Deref;
 use std::sync::Arc;
 
 use crate::simplify_expressions::ExprSimplifier;
-use crate::utils::collect_subquery_cols;
 
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter,
 };
 use datafusion_common::{plan_err, Column, DFSchemaRef, HashMap, Result, ScalarValue};
 use datafusion_expr::expr::Alias;
 use datafusion_expr::simplify::SimplifyContext;
-use datafusion_expr::utils::{conjunction, find_join_exprs, split_conjunction};
+use datafusion_expr::utils::{
+    collect_subquery_cols, conjunction, find_join_exprs, split_conjunction,
+};
 use datafusion_expr::{
     expr, lit, BinaryExpr, Cast, EmptyRelation, Expr, FetchType, LogicalPlan,
     LogicalPlanBuilder, Operator,

datafusion/optimizer/src/utils.rs

@@ -87,23 +87,6 @@ pub(crate) fn has_all_column_refs(expr: &Expr, schema_cols: &HashSet<Column>) ->
         == column_refs.len()
 }
 
-pub(crate) fn collect_subquery_cols(
-    exprs: &[Expr],
-    subquery_schema: &DFSchema,
-) -> Result<BTreeSet<Column>> {
-    exprs.iter().try_fold(BTreeSet::new(), |mut cols, expr| {
-        let mut using_cols: Vec<Column> = vec![];
-        for col in expr.column_refs().into_iter() {
-            if subquery_schema.has_column(col) {
-                using_cols.push(col.clone());
-            }
-        }
-
-        cols.extend(using_cols);
-        Result::<_>::Ok(cols)
-    })
-}
-
 pub(crate) fn replace_qualified_name(
     expr: Expr,
     cols: &BTreeSet<Column>,