diff --git a/datafusion/expr-common/src/placement.rs b/datafusion/expr-common/src/placement.rs index 8212ba618e322..f9fad423d97af 100644 --- a/datafusion/expr-common/src/placement.rs +++ b/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 + } + } } diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index 3428395b54222..b5aada9e8edeb 100644 --- a/datafusion/expr/src/expr.rs +++ b/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, } } diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs index 8a4abfcb48e56..89a5f01efde46 100644 --- a/datafusion/optimizer/src/extract_leaf_expressions.rs +++ b/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)? { diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 8a1dcc12ef874..b68522f96e3f4 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/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>, ) -> Result { + Ok(replace_cols_by_name_impl(e, replace_map)?.data) +} + +pub(super) fn replace_cols_by_name_impl( + e: Expr, + replace_map: &HashMap>, +) -> Result> { 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()?;