Avoid duplicating keep-in-place expressions

datafusion/expr-common/src/placement.rs

@@ -59,4 +59,31 @@ impl ExpressionPlacement {
             ExpressionPlacement::Column | ExpressionPlacement::MoveTowardsLeafNodes
         )
     }
+
+    pub fn reduce(placements: &[ExpressionPlacement]) -> ExpressionPlacement {
+        let mut all_literal = true;
+        let mut all_column = true;
+
+        for placement in placements {
+            match placement {
+                ExpressionPlacement::Literal => all_column = false,
+                ExpressionPlacement::Column => all_literal = false,
+                ExpressionPlacement::MoveTowardsLeafNodes => {
+                    all_column = false;
+                    all_literal = false;
+                }
+                ExpressionPlacement::KeepInPlace => {
+                    return ExpressionPlacement::KeepInPlace;
+                }
+            };
+        }
+
+        if all_literal {
+            ExpressionPlacement::Literal
+        } else if all_column {
+            ExpressionPlacement::Column
+        } else {
+            ExpressionPlacement::MoveTowardsLeafNodes
+        }
+    }
 }

datafusion/expr/src/expr.rs

@@ -60,6 +60,7 @@ pub use crate::sql::{
 // Moved in 51.0.0 to datafusion_common
 pub use datafusion_common::metadata::FieldMetadata;
 use datafusion_common::metadata::ScalarAndMetadata;
+use datafusion_expr_common::ExpressionPlacement::KeepInPlace;
 
 // This mirrors sqlparser::ast::NullTreatment but we need our own variant
 // for when the sql feature is disabled.
@@ -1654,13 +1655,35 @@ impl Expr {
         match self {
             Expr::Column(_) => ExpressionPlacement::Column,
             Expr::Literal(_, _) => ExpressionPlacement::Literal,
+            Expr::ScalarVariable(_, _) => ExpressionPlacement::Literal,
             Expr::Alias(inner) => inner.expr.placement(),
             Expr::ScalarFunction(func) => {
                 let arg_placements: Vec<_> =
                     func.args.iter().map(|arg| arg.placement()).collect();
                 func.func.placement(&arg_placements)
             }
-            _ => ExpressionPlacement::KeepInPlace,
+            Expr::BinaryExpr(BinaryExpr { left, right, .. }) => {
+                ExpressionPlacement::reduce(&[left.placement(), right.placement()])
+            }
+            Expr::Between(Between {
+                expr, low, high, ..
+            }) => ExpressionPlacement::reduce(&[
+                expr.placement(),
+                low.placement(),
+                high.placement(),
+            ]),
+            Expr::Not(e)
+            | Expr::IsNotNull(e)
+            | Expr::IsNull(e)
+            | Expr::IsTrue(e)
+            | Expr::IsFalse(e)
+            | Expr::IsUnknown(e)
+            | Expr::IsNotTrue(e)
+            | Expr::IsNotFalse(e)
+            | Expr::IsNotUnknown(e)
+            | Expr::Negative(e)
+            | Expr::Cast(Cast { expr: e, .. }) => e.placement(),
+            _ => KeepInPlace,
         }
     }
 

datafusion/optimizer/src/extract_leaf_expressions.rs

@@ -31,7 +31,7 @@ use datafusion_expr::logical_plan::LogicalPlan;
 use datafusion_expr::{Expr, ExpressionPlacement, Projection};
 
 use crate::optimizer::ApplyOrder;
-use crate::push_down_filter::replace_cols_by_name;
+use crate::push_down_filter::{replace_cols_by_name, replace_cols_by_name_impl};
 use crate::utils::{ColumnReference, has_all_column_refs, schema_columns};
 use crate::{OptimizerConfig, OptimizerRule};
 
@@ -672,17 +672,34 @@ fn build_extraction_projection_impl(
             })
             .collect();
 
+        let mut deferred_extractions = vec![];
+
         // Resolve column references through the projection's rename mapping
         let replace_map = build_projection_replace_map(existing);
 
         // Add new extracted expressions, resolving column refs through the projection
         for (expr, alias) in extracted_exprs {
-            let resolved = replace_cols_by_name(expr.clone().alias(alias), &replace_map)?;
+            let Transformed {
+                data: resolved,
+                transformed,
+                ..
+            } = replace_cols_by_name_impl(expr.clone().alias(alias), &replace_map)?;
             let resolved_inner = if let Expr::Alias(a) = &resolved {
                 a.expr.as_ref()
             } else {
                 &resolved
             };
+
+            // If any columns were inline and the end result is keep-in-place we're likely to be
+            // duplicating an expensive expression. Defer the extraction to a second projection to
+            // avoid this.
+            if transformed
+                && resolved_inner.placement() == ExpressionPlacement::KeepInPlace
+            {
+                deferred_extractions.push(expr.clone().alias(alias));
+                continue;
+            }
+
             if let Some(existing_alias) = existing_extractions.get(resolved_inner) {
                 // Same expression already extracted under a different alias —
                 // add the expression with the new alias so both names are
@@ -727,9 +744,23 @@ fn build_extraction_projection_impl(
             // If resolved to non-column expr, it's already computed by existing projection
         }
 
-        Projection::try_new(proj_exprs, Arc::clone(&existing.input))
+        let extended_projection =
+            Projection::try_new(proj_exprs, Arc::clone(&existing.input))?;
+
+        if deferred_extractions.is_empty() {
+            Ok(extended_projection)
+        } else {
+            let mut proj_exprs = Vec::new();
+            proj_exprs.extend(deferred_extractions);
+            for (qualifier, field) in extended_projection.schema.as_ref().iter() {
+                proj_exprs.push(Expr::from((qualifier, field)));
+            }
+            Projection::try_new(
+                proj_exprs,
+                Arc::new(LogicalPlan::Projection(extended_projection)),
+            )
+        }
     } else {
-        // Build new projection with extracted expressions + all input columns
         let mut proj_exprs = Vec::new();
         for (expr, alias) in extracted_exprs {
             proj_exprs.push(expr.clone().alias(alias));
@@ -1191,6 +1222,9 @@ fn push_extraction_pairs(
             // the (None, true) fallback can't find the original aliases.
             // This handles: Extraction → Recovery(cols) → Filter → ... → TableScan
             // by pushing through the recovery projection AND the filter in one pass.
+
+            // TODO when `build_extraction_projection_impl` defers extraction, `try_push_input`
+            // will lead to an infinite recursion
             if is_pure_extraction_projection(&merged_plan)
                 && let Some(pushed) = try_push_input(&merged_plan, alias_generator)?
             {

datafusion/optimizer/src/push_down_filter.rs

@@ -26,9 +26,7 @@ use itertools::Itertools;
 use log::{Level, debug, log_enabled};
 
 use datafusion_common::instant::Instant;
-use datafusion_common::tree_node::{
-    Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
-};
+use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
 use datafusion_common::{
     Column, DFSchema, Result, assert_eq_or_internal_err, internal_err, plan_err,
     qualified_name,
@@ -1312,6 +1310,7 @@ fn rewrite_projection(
         .partition(|(_, value)| {
             value.is_volatile()
                 || value.placement() == ExpressionPlacement::MoveTowardsLeafNodes
+                || value.placement() == ExpressionPlacement::KeepInPlace
         });
 
     let mut push_predicates = vec![];
@@ -1376,6 +1375,13 @@ pub fn replace_cols_by_name(
     e: Expr,
     replace_map: &HashMap<String, impl AsRef<Expr>>,
 ) -> Result<Expr> {
+    Ok(replace_cols_by_name_impl(e, replace_map)?.data)
+}
+
+pub(super) fn replace_cols_by_name_impl(
+    e: Expr,
+    replace_map: &HashMap<String, impl AsRef<Expr>>,
+) -> Result<Transformed<Expr>> {
     e.transform_up(|expr| {
         if let Expr::Column(c) = &expr
             && let Some(new_expr) = replace_map.get(&c.flat_name())
@@ -1385,7 +1391,6 @@ pub fn replace_cols_by_name(
             Ok(Transformed::no(expr))
         }
     })
-    .data()
 }
 
 /// Unalias expression reference.
@@ -1454,7 +1459,7 @@ mod tests {
     use crate::assert_optimized_plan_eq_snapshot;
     use crate::optimizer::Optimizer;
     use crate::simplify_expressions::SimplifyExpressions;
-    use crate::test::udfs::leaf_udf_expr;
+    use crate::test::udfs::{PlacementTestUDF, get_field_like, leaf_udf_expr};
     use crate::test::*;
     use datafusion_expr::test::function_stub::sum;
     use insta::assert_snapshot;
@@ -4394,6 +4399,37 @@ mod tests {
         )
     }
 
+    #[test]
+    fn filter_not_pushed_through_nested_computed_projection() -> Result<()> {
+        let udf = ScalarUDF::new_from_impl(
+            PlacementTestUDF::new().with_placement(ExpressionPlacement::KeepInPlace),
+        );
+        let inner = LogicalPlanBuilder::from(test_table_scan()?)
+            .project(vec![udf.call(vec![col("a")]).alias("c1"), col("b")])?
+            .build()?;
+        let outer = LogicalPlanBuilder::from(inner)
+            .project(vec![
+                get_field_like(col("c1"), "x").alias("c2"),
+                col("c1"),
+                col("b"),
+            ])?
+            .build()?;
+        let plan = LogicalPlanBuilder::from(outer)
+            .filter(col("c1").is_not_null().and(col("c2").is_null()))?
+            .build()?;
+
+        assert_optimized_plan_equal!(
+            plan,
+            @r#"
+        Filter: c2 IS NULL
+          Projection: get_field_like(c1, Utf8("x")) AS c2, c1, test.b
+            Filter: c1 IS NOT NULL
+              Projection: keep_in_place_udf(test.a) AS c1, test.b
+                TableScan: test
+        "#
+        )
+    }
+
     #[test]
     fn filter_not_pushed_down_through_table_scan_with_fetch() -> Result<()> {
         let scan = test_table_scan()?;