Skip to content

Commit dbf62c1

Browse files
authored
fix: Add a subquery when unparsing a Filter depending on a Projection (#25549)
## Which issue does this PR close? - Closes #25548. ## Rationale for this change When we have a `Filter` next to a `Projection`, we might not be able to unparse them at the same level if the predicate depends on columns aliased by the projection. For those cases, we need to add a subquery in between. ## What changes are included in this PR? - Iff any column of a `Filter` is aliased by a `Projection` directly below it, a subquery is added with `derive`. ## What is the testing strategy for this PR? Unit tests. ## Are there any user-facing changes? No, this only adds the subquery when it is necessary, so existing tests remain unchanged.
1 parent bdad988 commit dbf62c1

3 files changed

Lines changed: 175 additions & 7 deletions

File tree

‎datafusion/core/tests/sql/unparser.rs‎

Lines changed: 124 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ use datafusion_catalog::memory::MemorySchemaProvider;
5151
use datafusion_catalog::{CatalogProvider, MemoryCatalogProvider, SchemaProvider};
5252
use datafusion_common::Column;
5353
use datafusion_expr::Expr;
54-
use datafusion_sql::unparser::dialect::{DefaultDialect, DuckDBDialect};
54+
use datafusion_sql::unparser::dialect::{DefaultDialect, DuckDBDialect, MySqlDialect};
5555
use datafusion_sql::unparser::{Unparser, plan_to_sql};
5656
use itertools::Itertools;
5757
use recursive::{set_minimum_stack_size, set_stack_allocation_size};
@@ -772,6 +772,129 @@ async fn optimized_filter_with_subquery_alias() -> Result<()> {
772772
Ok(())
773773
}
774774

775+
#[tokio::test]
776+
async fn optimized_filter_after_projection() -> Result<()> {
777+
let ctx = SessionContext::new();
778+
ctx.sql("create table t (a bigint)")
779+
.await?
780+
.collect()
781+
.await?;
782+
783+
// x=1 cannot be pushed to the inner subquery since it depends on the projection
784+
let df = ctx
785+
.sql(
786+
"
787+
select *
788+
from (
789+
select random() as x
790+
from t
791+
)
792+
where x = 1
793+
",
794+
)
795+
.await?;
796+
let plan = df.into_optimized_plan()?;
797+
let sql = plan_to_sql(&plan)?.to_string();
798+
assert_eq!(
799+
sql,
800+
"SELECT * FROM (SELECT random() AS x FROM t) WHERE (x = 1.0)"
801+
);
802+
803+
// a=1 can be pushed since it does not depend on the projection
804+
let df = ctx
805+
.sql(
806+
"
807+
select *
808+
from (
809+
select a
810+
from t
811+
)
812+
where a = 1
813+
",
814+
)
815+
.await?;
816+
let plan = df.into_optimized_plan()?;
817+
let sql = plan_to_sql(&plan)?.to_string();
818+
assert_eq!(sql, "SELECT t.a FROM t WHERE (t.a = 1)");
819+
820+
// a=1 is pushed but x=1 is not
821+
let df = ctx
822+
.sql(
823+
"
824+
select *
825+
from (
826+
select a, random() as x
827+
from t
828+
)
829+
where a = 1 and x = 1
830+
",
831+
)
832+
.await?;
833+
let plan = df.into_optimized_plan()?;
834+
let sql = plan_to_sql(&plan)?.to_string();
835+
assert_eq!(
836+
sql,
837+
"SELECT * FROM (SELECT t.a, random() AS x FROM t WHERE (t.a = 1)) WHERE (x = 1.0)"
838+
);
839+
840+
// b=1 is optimized into a+1=1 and pushed down
841+
let df = ctx
842+
.sql(
843+
"
844+
select *
845+
from (
846+
select a + 1 as b
847+
from t
848+
)
849+
where b = 1
850+
",
851+
)
852+
.await?;
853+
let plan = df.into_optimized_plan()?;
854+
let sql = plan_to_sql(&plan)?.to_string();
855+
assert_eq!(sql, "SELECT (t.a + 1) AS b FROM t WHERE ((t.a + 1) = 1)");
856+
857+
// a+1=1 is also optimized and pushed down
858+
let df = ctx
859+
.sql(
860+
"
861+
select *
862+
from (
863+
select a + 1 as a
864+
from t
865+
)
866+
where a + 1 = 1
867+
",
868+
)
869+
.await?;
870+
let plan = df.into_optimized_plan()?;
871+
let sql = plan_to_sql(&plan)?.to_string();
872+
assert_eq!(
873+
sql,
874+
"SELECT (t.a + 1) AS a FROM t WHERE (((t.a + 1) + 1) = 1)"
875+
);
876+
877+
// x>a adds a subquery but converts "t.a" to "a"
878+
let df = ctx
879+
.sql("SELECT * FROM (SELECT a, random() AS x FROM t) WHERE x > a")
880+
.await?;
881+
let plan = df.into_optimized_plan()?;
882+
let sql = plan_to_sql(&plan)?.to_string();
883+
assert_eq!(
884+
sql,
885+
"SELECT * FROM (SELECT t.a, random() AS x FROM t) WHERE (x > CAST(a AS DOUBLE))"
886+
);
887+
// dialects like MySQL require aliasing subqueries
888+
let unparser = Unparser::new(&MySqlDialect {});
889+
let sql = unparser.plan_to_sql(&plan)?.to_string();
890+
assert_eq!(
891+
sql,
892+
"SELECT * FROM (SELECT `t`.`a`, random() AS `x` FROM `t`) AS `derived_projection` WHERE (`x` > CAST(`derived_projection`.`a` AS DOUBLE))"
893+
);
894+
895+
Ok(())
896+
}
897+
775898
/// The outcome of running a single roundtrip test.
776899
///
777900
/// A successful test produces [`TestCaseResult::Success`].

‎datafusion/sql/src/unparser/plan.rs‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,13 +33,14 @@ use super::{
3333
unproject_unnest_expr_as_flatten_value, unproject_window_exprs,
3434
},
3535
};
36-
use crate::unparser::extension_unparser::{
37-
UnparseToStatementResult, UnparseWithinStatementResult,
38-
};
3936
use crate::unparser::utils::{find_unnest_node_until_relation, unproject_agg_exprs};
4037
use crate::unparser::{
4138
ast::FlattenRelationBuilder, ast::UnnestRelationBuilder, rewrite::rewrite_qualify,
4239
};
40+
use crate::unparser::{
41+
extension_unparser::{UnparseToStatementResult, UnparseWithinStatementResult},
42+
utils::filter_depends_on_input_alias,
43+
};
4344
use crate::utils::UNNEST_PLACEHOLDER;
4445
use datafusion_common::{
4546
Column, DFSchema, DataFusionError, Result, ScalarValue, TableReference,
@@ -518,6 +519,14 @@ impl Unparser<'_> {
518519
schema: projection.input.schema().as_ref(),
519520
})
520521
}
522+
// when a filter references aliases from its input, we add a subquery around it to
523+
// prevent invalid references, so we need to keep track of it here
524+
LogicalPlan::Filter(filter) if filter_depends_on_input_alias(filter) => {
525+
Some(DerivedInputScope {
526+
alias: "derived_projection",
527+
schema: filter.input.schema().as_ref(),
528+
})
529+
}
521530
LogicalPlan::Filter(filter) => {
522531
Self::derived_input_scope(filter.input.as_ref(), select)
523532
}
@@ -1243,6 +1252,18 @@ impl Unparser<'_> {
12431252
select.selection(Some(filter_expr));
12441253
}
12451254

1255+
// if the inner plan aliases columns used by the filter, we need to convert to a
1256+
// subquery to prevent invalid references
1257+
if filter_depends_on_input_alias(filter) {
1258+
return self.derive_with_dialect_alias(
1259+
"derived_projection",
1260+
&filter.input,
1261+
relation,
1262+
false,
1263+
vec![],
1264+
);
1265+
}
1266+
12461267
self.select_to_sql_recursively(
12471268
filter.input.as_ref(),
12481269
query,

‎datafusion/sql/src/unparser/utils.rs‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use std::{cmp::Ordering, str::FromStr, sync::Arc, vec};
18+
use std::{cmp::Ordering, collections::HashSet, str::FromStr, sync::Arc, vec};
1919

2020
use super::{
2121
Unparser, dialect::CharacterLengthStyle, dialect::DateFieldExtractStyle,
@@ -27,8 +27,8 @@ use datafusion_common::{
2727
tree_node::{Transformed, TransformedResult, TreeNode},
2828
};
2929
use datafusion_expr::{
30-
Aggregate, Expr, LogicalPlan, LogicalPlanBuilder, Projection, SortExpr, Unnest,
31-
Window, expr, utils::grouping_set_to_exprlist,
30+
Aggregate, Expr, Filter, LogicalPlan, LogicalPlanBuilder, Projection, SortExpr,
31+
Unnest, Window, expr, utils::grouping_set_to_exprlist,
3232
};
3333

3434
use arrow::compute::DatePart;
@@ -609,3 +609,27 @@ pub(crate) fn sqlite_date_trunc_to_sql(
609609

610610
Ok(None)
611611
}
612+
613+
/// Returns true if any column used by a `Filter` depends on aliased expressions by its input. This
614+
/// should only happen when we have a Filter + Projection plan.
615+
pub(crate) fn filter_depends_on_input_alias(filter: &Filter) -> bool {
616+
match filter.input.as_ref() {
617+
LogicalPlan::Projection(projection) => {
618+
let alias = projection
619+
.expr
620+
.iter()
621+
.filter_map(|e| match e {
622+
Expr::Alias(a) => Some((&a.relation, &a.name)),
623+
_ => None,
624+
})
625+
.collect::<HashSet<_>>();
626+
627+
filter
628+
.predicate
629+
.column_refs()
630+
.iter()
631+
.any(|c| alias.contains(&(&c.relation, &c.name)))
632+
}
633+
_ => false,
634+
}
635+
}

0 commit comments

Comments
 (0)