fix: Add a subquery when unparsing a Filter depending on a Projection

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`].

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,

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::<HashSet<_>>();
+
+            filter
+                .predicate
+                .column_refs()
+                .iter()
+                .any(|c| alias.contains(&(&c.relation, &c.name)))
+        }
+        _ => false,
+    }
+}