Avoid allocating column names

Author: AdamGS

datafusion/optimizer/src/extract_leaf_expressions.rs

@@ -32,7 +32,7 @@ use datafusion_expr::{Expr, ExpressionPlacement, Projection};
 
 use crate::optimizer::ApplyOrder;
 use crate::push_down_filter::replace_cols_by_name;
-use crate::utils::has_all_column_refs;
+use crate::utils::{ColumnReference, has_all_column_refs, schema_columns};
 use crate::{OptimizerConfig, OptimizerRule};
 
 /// Prefix for aliases generated by the extraction optimizer passes.
@@ -213,10 +213,11 @@ fn extract_from_plan(
         .collect();
 
     // Build per-input column sets for routing expressions to the correct input
-    let input_column_sets: Vec<std::collections::HashSet<Column>> = input_schemas
-        .iter()
-        .map(|schema| schema_columns(schema.as_ref()))
-        .collect();
+    let input_column_sets: Vec<std::collections::HashSet<ColumnReference>> =
+        input_schemas
+            .iter()
+            .map(|schema| schema_columns(schema.as_ref()))
+            .collect();
 
     // Transform expressions via map_expressions with routing
     let transformed = plan.map_expressions(|expr| {
@@ -272,7 +273,7 @@ fn extract_from_plan(
 /// in both sides of a join).
 fn find_owning_input(
     expr: &Expr,
-    input_column_sets: &[std::collections::HashSet<Column>],
+    input_column_sets: &[std::collections::HashSet<ColumnReference>],
 ) -> Option<usize> {
     let mut found = None;
     for (idx, cols) in input_column_sets.iter().enumerate() {
@@ -292,7 +293,7 @@ fn find_owning_input(
 fn routing_extract(
     expr: Expr,
     extractors: &mut [LeafExpressionExtractor],
-    input_column_sets: &[std::collections::HashSet<Column>],
+    input_column_sets: &[std::collections::HashSet<ColumnReference>],
 ) -> Result<Transformed<Expr>> {
     expr.transform_down(|e| {
         // Skip expressions already aliased with extracted expression pattern
@@ -340,19 +341,6 @@ fn routing_extract(
     })
 }
 
-/// Returns all columns in the schema (both qualified and unqualified forms)
-fn schema_columns(schema: &DFSchema) -> std::collections::HashSet<Column> {
-    schema
-        .iter()
-        .flat_map(|(qualifier, field)| {
-            [
-                Column::new(qualifier.cloned(), field.name()),
-                Column::new_unqualified(field.name()),
-            ]
-        })
-        .collect()
-}
-
 /// Rewrites extraction pairs and column references from one qualifier
 /// space to another.
 ///
@@ -1072,7 +1060,7 @@ fn route_to_inputs(
     pairs: &[(Expr, String)],
     columns: &IndexSet<Column>,
     node: &LogicalPlan,
-    input_column_sets: &[std::collections::HashSet<Column>],
+    input_column_sets: &[std::collections::HashSet<ColumnReference>],
     input_schemas: &[Arc<DFSchema>],
 ) -> Result<Option<Vec<ExtractionTarget>>> {
     let num_inputs = input_schemas.len();
@@ -1173,7 +1161,7 @@ fn try_push_into_inputs(
     // Build per-input schemas and column sets for routing
     let input_schemas: Vec<Arc<DFSchema>> =
         inputs.iter().map(|i| Arc::clone(i.schema())).collect();
-    let input_column_sets: Vec<std::collections::HashSet<Column>> =
+    let input_column_sets: Vec<std::collections::HashSet<ColumnReference>> =
         input_schemas.iter().map(|s| schema_columns(s)).collect();
 
     // Route pairs and columns to the appropriate inputs
@@ -2436,16 +2424,18 @@ mod tests {
         // Simulate schema_columns output for two sides of a join where both
         // have a "user" column — each set contains the qualified and
         // unqualified form.
-        let left_cols: HashSet<Column> = [
-            Column::new(Some("test"), "user"),
-            Column::new_unqualified("user"),
+        let relation = "test".into();
+        let left_cols: HashSet<ColumnReference> = [
+            ColumnReference::new(Some(&relation), "user"),
+            ColumnReference::new_unqualified("user"),
         ]
         .into_iter()
         .collect();
 
-        let right_cols: HashSet<Column> = [
-            Column::new(Some("right"), "user"),
-            Column::new_unqualified("user"),
+        let relation = "right".into();
+        let right_cols: HashSet<ColumnReference> = [
+            ColumnReference::new(Some(&relation), "user"),
+            ColumnReference::new_unqualified("user"),
         ]
         .into_iter()
         .collect();

datafusion/optimizer/src/push_down_filter.rs

@@ -45,7 +45,9 @@ use datafusion_expr::{
 
 use crate::optimizer::ApplyOrder;
 use crate::simplify_expressions::simplify_predicates;
-use crate::utils::{has_all_column_refs, is_restrict_null_predicate};
+use crate::utils::{
+    ColumnReference, has_all_column_refs, is_restrict_null_predicate, schema_columns,
+};
 use crate::{OptimizerConfig, OptimizerRule};
 use datafusion_expr::ExpressionPlacement;
 
@@ -190,11 +192,11 @@ struct ColumnChecker<'a> {
     /// schema of left join input
     left_schema: &'a DFSchema,
     /// columns in left_schema, computed on demand
-    left_columns: Option<HashSet<Column>>,
+    left_columns: Option<HashSet<ColumnReference<'a>>>,
     /// schema of right join input
     right_schema: &'a DFSchema,
     /// columns in left_schema, computed on demand
-    right_columns: Option<HashSet<Column>>,
+    right_columns: Option<HashSet<ColumnReference<'a>>>,
 }
 
 impl<'a> ColumnChecker<'a> {
@@ -224,20 +226,6 @@ impl<'a> ColumnChecker<'a> {
     }
 }
 
-/// Returns all columns in the schema
-fn schema_columns(schema: &DFSchema) -> HashSet<Column> {
-    schema
-        .iter()
-        .flat_map(|(qualifier, field)| {
-            [
-                Column::new(qualifier.cloned(), field.name()),
-                // we need to push down filter using unqualified column as well
-                Column::new_unqualified(field.name()),
-            ]
-        })
-        .collect::<HashSet<_>>()
-}
-
 /// Determine whether the predicate can evaluate as the join conditions
 fn can_evaluate_as_join_condition(predicate: &Expr) -> Result<bool> {
     let mut is_evaluate = true;
@@ -320,7 +308,7 @@ fn can_evaluate_as_join_condition(predicate: &Expr) -> Result<bool> {
 /// * do nothing.
 fn extract_or_clauses_for_join<'a>(
     filters: &'a [Expr],
-    schema_cols: &'a HashSet<Column>,
+    schema_cols: &'a HashSet<ColumnReference>,
 ) -> impl Iterator<Item = Expr> + 'a {
     // new formed OR clauses and their column references
     filters.iter().filter_map(move |expr| {
@@ -353,7 +341,10 @@ fn extract_or_clauses_for_join<'a>(
 /// Otherwise, return None.
 ///
 /// For other clause, apply the rule above to extract clause.
-fn extract_or_clause(expr: &Expr, schema_columns: &HashSet<Column>) -> Option<Expr> {
+fn extract_or_clause(
+    expr: &Expr,
+    schema_columns: &HashSet<ColumnReference>,
+) -> Option<Expr> {
     let mut predicate = None;
 
     match expr {

datafusion/optimizer/src/utils.rs

@@ -22,6 +22,7 @@ use std::collections::{BTreeSet, HashMap, HashSet};
 use crate::analyzer::type_coercion::TypeCoercionRewriter;
 use arrow::array::{Array, RecordBatch, new_null_array};
 use arrow::datatypes::{DataType, Field, Schema};
+use datafusion_common::TableReference;
 use datafusion_common::cast::as_boolean_array;
 use datafusion_common::tree_node::{TransformedResult, TreeNode};
 use datafusion_common::{Column, DFSchema, Result, ScalarValue};
@@ -37,12 +38,17 @@ use std::sync::Arc;
 pub use datafusion_expr::expr_rewriter::NamePreserver;
 
 /// Returns true if `expr` contains all columns in `schema_cols`
-pub(crate) fn has_all_column_refs(expr: &Expr, schema_cols: &HashSet<Column>) -> bool {
+pub(crate) fn has_all_column_refs(
+    expr: &Expr,
+    schema_cols: &HashSet<ColumnReference<'_>>,
+) -> bool {
     let column_refs = expr.column_refs();
     // note can't use HashSet::intersect because of different types (owned vs References)
-    schema_cols
+    column_refs
         .iter()
-        .filter(|c| column_refs.contains(c))
+        .filter(|c| {
+            schema_cols.contains(&ColumnReference::new(c.relation.as_ref(), c.name()))
+        })
         .count()
         == column_refs.len()
 }
@@ -62,6 +68,40 @@ pub(crate) fn replace_qualified_name(
     replace_col(expr, &replace_map)
 }
 
+///Column reference to avoid copying string around
+#[derive(PartialEq, Eq, Hash, Debug)]
+pub(crate) struct ColumnReference<'a> {
+    pub relation: Option<&'a TableReference>,
+    pub name: &'a str,
+}
+
+impl<'a> ColumnReference<'a> {
+    pub fn new(relation: Option<&'a TableReference>, name: &'a str) -> Self {
+        Self { relation, name }
+    }
+
+    pub fn new_unqualified(name: &'a str) -> Self {
+        Self {
+            relation: None,
+            name,
+        }
+    }
+}
+
+/// Returns references to all columns in the schema
+pub(crate) fn schema_columns(schema: &DFSchema) -> HashSet<ColumnReference<'_>> {
+    schema
+        .iter()
+        .flat_map(|(qualifier, field)| {
+            [
+                ColumnReference::new(qualifier, field.name()),
+                // we need to push down filter using unqualified column as well
+                ColumnReference::new_unqualified(field.name()),
+            ]
+        })
+        .collect::<HashSet<_>>()
+}
+
 /// Log the plan in debug/tracing mode after some part of the optimizer runs
 pub fn log_plan(description: &str, plan: &LogicalPlan) {
     debug!("{description}:\n{}\n", plan.display_indent());