diff --git a/datafusion/core/tests/sql/unparser.rs b/datafusion/core/tests/sql/unparser.rs index 3982c60dbc7d9..4aa7f0890e8ec 100644 --- a/datafusion/core/tests/sql/unparser.rs +++ b/datafusion/core/tests/sql/unparser.rs @@ -51,8 +51,8 @@ use datafusion_catalog::memory::MemorySchemaProvider; use datafusion_catalog::{CatalogProvider, MemoryCatalogProvider, SchemaProvider}; use datafusion_common::Column; use datafusion_expr::Expr; -use datafusion_sql::unparser::Unparser; use datafusion_sql::unparser::dialect::{DefaultDialect, DuckDBDialect}; +use datafusion_sql::unparser::{Unparser, plan_to_sql}; use itertools::Itertools; use recursive::{set_minimum_stack_size, set_stack_allocation_size}; @@ -747,6 +747,111 @@ async fn optimized_duckdb_unparse_top_level_sort_over_agg_uses_select_alias() -> Ok(()) } +#[tokio::test] +async fn optimized_filter_after_projection() -> Result<()> { + let ctx = SessionContext::new(); + ctx.sql("create table t (a bigint)") + .await? + .collect() + .await?; + + // x=1 cannot be pushed to the inner subquery since it depends on the projection + let df = ctx + .sql( + " + select * + from ( + select random() as x + from t + ) + where x = 1 + ", + ) + .await?; + let plan = df.into_optimized_plan()?; + let sql = plan_to_sql(&plan)?.to_string(); + assert_eq!( + sql, + "SELECT * FROM (SELECT random() AS x FROM t) WHERE (x = 1.0)" + ); + + // a=1 can be pushed since it does not depend on the projection + let df = ctx + .sql( + " + select * + from ( + select a + from t + ) + where a = 1 + ", + ) + .await?; + let plan = df.into_optimized_plan()?; + let sql = plan_to_sql(&plan)?.to_string(); + assert_eq!(sql, "SELECT t.a FROM t WHERE (t.a = 1)"); + + // a=1 is pushed but x=1 is not + let df = ctx + .sql( + " + select * + from ( + select a, random() as x + from t + ) + where a = 1 and x = 1 + ", + ) + .await?; + let plan = df.into_optimized_plan()?; + let sql = plan_to_sql(&plan)?.to_string(); + assert_eq!( + sql, + "SELECT * FROM (SELECT t.a, random() AS x FROM t WHERE (t.a = 1)) WHERE (x = 1.0)" + ); + + // b=1 is optimized into a+1=1 and pushed down + let df = ctx + .sql( + " + select * + from ( + select a + 1 as b + from t + ) + where b = 1 + ", + ) + .await?; + let plan = df.into_optimized_plan()?; + let sql = plan_to_sql(&plan)?.to_string(); + assert_eq!(sql, "SELECT (t.a + 1) AS b FROM t WHERE ((t.a + 1) = 1)"); + + // a+1=1 is also optimized and pushed down + let df = ctx + .sql( + " + select * + from ( + select a + 1 as a + from t + ) + where a + 1 = 1 + ", + ) + .await?; + let plan = df.into_optimized_plan()?; + let sql = plan_to_sql(&plan)?.to_string(); + assert_eq!( + sql, + "SELECT (t.a + 1) AS a FROM t WHERE (((t.a + 1) + 1) = 1)" + ); + + Ok(()) +} + /// The outcome of running a single roundtrip test. /// /// A successful test produces [`TestCaseResult::Success`]. diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 18af08fc18361..eb770ba25bcfb 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -33,13 +33,14 @@ use super::{ unproject_unnest_expr_as_flatten_value, unproject_window_exprs, }, }; -use crate::unparser::extension_unparser::{ - UnparseToStatementResult, UnparseWithinStatementResult, -}; use crate::unparser::utils::{find_unnest_node_until_relation, unproject_agg_exprs}; use crate::unparser::{ ast::FlattenRelationBuilder, ast::UnnestRelationBuilder, rewrite::rewrite_qualify, }; +use crate::unparser::{ + extension_unparser::{UnparseToStatementResult, UnparseWithinStatementResult}, + utils::filter_depends_on_input_alias, +}; use crate::utils::UNNEST_PLACEHOLDER; use datafusion_common::{ Column, DFSchema, DataFusionError, Result, ScalarValue, TableReference, @@ -1232,6 +1233,12 @@ impl Unparser<'_> { select.selection(Some(filter_expr)); } + // if the inner plan aliases columns used by the filter, we need to convert to a + // subquery to prevent invalid references + if filter_depends_on_input_alias(filter) { + return self.derive(&filter.input, relation, None, false); + } + self.select_to_sql_recursively( filter.input.as_ref(), query, diff --git a/datafusion/sql/src/unparser/utils.rs b/datafusion/sql/src/unparser/utils.rs index 86f3e23115dcb..8179de63b4e97 100644 --- a/datafusion/sql/src/unparser/utils.rs +++ b/datafusion/sql/src/unparser/utils.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::{cmp::Ordering, str::FromStr, sync::Arc, vec}; +use std::{cmp::Ordering, collections::HashSet, str::FromStr, sync::Arc, vec}; use super::{ Unparser, dialect::CharacterLengthStyle, dialect::DateFieldExtractStyle, @@ -27,8 +27,8 @@ use datafusion_common::{ tree_node::{Transformed, TransformedResult, TreeNode}, }; use datafusion_expr::{ - Aggregate, Expr, LogicalPlan, LogicalPlanBuilder, Projection, SortExpr, Unnest, - Window, expr, utils::grouping_set_to_exprlist, + Aggregate, Expr, Filter, LogicalPlan, LogicalPlanBuilder, Projection, SortExpr, + Unnest, Window, expr, utils::grouping_set_to_exprlist, }; use arrow::compute::DatePart; @@ -609,3 +609,27 @@ pub(crate) fn sqlite_date_trunc_to_sql( Ok(None) } + +/// Returns true if any column used by a `Filter` depends on aliased expressions by its input. This +/// should only happen when we have a Filter + Projection plan. +pub(crate) fn filter_depends_on_input_alias(filter: &Filter) -> bool { + match filter.input.as_ref() { + LogicalPlan::Projection(projection) => { + let alias = projection + .expr + .iter() + .filter_map(|e| match e { + Expr::Alias(a) => Some((&a.relation, &a.name)), + _ => None, + }) + .collect::>(); + + filter + .predicate + .column_refs() + .iter() + .any(|c| alias.contains(&(&c.relation, &c.name))) + } + _ => false, + } +}