From b7fa3c82d95fb685339379b9cd5dcbfe25f91cec Mon Sep 17 00:00:00 2001 From: Samuel Colvin Date: Sat, 6 Jul 2024 17:00:30 +0100 Subject: [PATCH 1/4] allow alias in predicate, fix #11306 --- .../user_defined/user_defined_sql_planner.rs | 37 ++++++++++++++++++- datafusion/expr/src/logical_plan/plan.rs | 11 +----- 2 files changed, 37 insertions(+), 11 deletions(-) diff --git a/datafusion/core/tests/user_defined/user_defined_sql_planner.rs b/datafusion/core/tests/user_defined/user_defined_sql_planner.rs index 37df7e0900b4b..50d19f2780609 100644 --- a/datafusion/core/tests/user_defined/user_defined_sql_planner.rs +++ b/datafusion/core/tests/user_defined/user_defined_sql_planner.rs @@ -24,6 +24,8 @@ use datafusion::execution::FunctionRegistry; use datafusion::logical_expr::Operator; use datafusion::prelude::*; use datafusion::sql::sqlparser::ast::BinaryOperator; +use datafusion_common::ScalarValue; +use datafusion_expr::expr::Alias; use datafusion_expr::planner::{PlannerResult, RawBinaryExpr, UserDefinedSQLPlanner}; use datafusion_expr::BinaryExpr; @@ -50,13 +52,22 @@ impl UserDefinedSQLPlanner for MyCustomPlanner { op: Operator::Plus, }))) } + BinaryOperator::Question => { + Ok(PlannerResult::Planned(Expr::Alias(Alias::new( + Expr::Literal(ScalarValue::Boolean(Some(true))), + None::<&str>, + format!("{} ? {}", expr.left, expr.right), + )))) + } _ => Ok(PlannerResult::Original(expr)), } } } async fn plan_and_collect(sql: &str) -> Result> { - let mut ctx = SessionContext::new(); + let config = + SessionConfig::new().set_str("datafusion.sql_parser.dialect", "postgres"); + let mut ctx = SessionContext::new_with_config(config); ctx.register_user_defined_sql_planner(Arc::new(MyCustomPlanner))?; ctx.sql(sql).await?.collect().await } @@ -86,3 +97,27 @@ async fn test_custom_operators_long_arrow() { ]; assert_batches_eq!(&expected, &actual); } + +#[tokio::test] +async fn test_question_sekect() { + let actual = plan_and_collect("select a ? 2 from (select 1 as a);") + .await + .unwrap(); + let expected = [ + "+--------------+", + "| a ? Int64(2) |", + "+--------------+", + "| true |", + "+--------------+", + ]; + assert_batches_eq!(&expected, &actual); +} + +#[tokio::test] +async fn test_question_filter() { + let actual = plan_and_collect("select a from (select 1 as a) where a ? 2;") + .await + .unwrap(); + let expected = ["+---+", "| a |", "+---+", "| 1 |", "+---+"]; + assert_batches_eq!(&expected, &actual); +} diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 8fd5982a0f2e4..7ca08e729eb86 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -25,7 +25,7 @@ use std::sync::Arc; use super::dml::CopyTo; use super::DdlStatement; use crate::builder::{change_redundant_column, unnest_with_options}; -use crate::expr::{Alias, Placeholder, Sort as SortExpr, WindowFunction}; +use crate::expr::{Placeholder, Sort as SortExpr, WindowFunction}; use crate::expr_rewriter::{create_col_from_scalar_expr, normalize_cols}; use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor}; use crate::logical_plan::extension::UserDefinedLogicalNode; @@ -2130,15 +2130,6 @@ impl Filter { } } - // filter predicates should not be aliased - if let Expr::Alias(Alias { expr, name, .. }) = predicate { - return plan_err!( - "Attempted to create Filter predicate with \ - expression `{expr}` aliased as '{name}'. Filter predicates should not be \ - aliased." - ); - } - Ok(Self { predicate, input }) } From 7f6759d8879cdcc51ad89c40d6d3cf03aab210bd Mon Sep 17 00:00:00 2001 From: Samuel Colvin Date: Sun, 7 Jul 2024 19:42:34 +0100 Subject: [PATCH 2/4] Fix typo. Co-authored-by: Andrew Lamb --- datafusion/core/tests/user_defined/user_defined_sql_planner.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/core/tests/user_defined/user_defined_sql_planner.rs b/datafusion/core/tests/user_defined/user_defined_sql_planner.rs index 50d19f2780609..6b660e15a1970 100644 --- a/datafusion/core/tests/user_defined/user_defined_sql_planner.rs +++ b/datafusion/core/tests/user_defined/user_defined_sql_planner.rs @@ -99,7 +99,7 @@ async fn test_custom_operators_long_arrow() { } #[tokio::test] -async fn test_question_sekect() { +async fn test_question_select() { let actual = plan_and_collect("select a ? 2 from (select 1 as a);") .await .unwrap(); From 542d9d150c06a2c1a445210753e48eb68345b2e6 Mon Sep 17 00:00:00 2001 From: Samuel Colvin Date: Mon, 8 Jul 2024 10:21:00 +0100 Subject: [PATCH 3/4] unalise predicate --- datafusion/expr/src/logical_plan/plan.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 7ca08e729eb86..d91c38b828442 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -2130,7 +2130,10 @@ impl Filter { } } - Ok(Self { predicate, input }) + Ok(Self { + predicate: predicate.unalias(), + input, + }) } /// Is this filter guaranteed to return 0 or 1 row in a given instantiation? From dc58b801b750ae44c858bb64eeb423e4b1d9eff5 Mon Sep 17 00:00:00 2001 From: Samuel Colvin Date: Mon, 8 Jul 2024 15:55:03 +0100 Subject: [PATCH 4/4] use unalias_nested --- datafusion/expr/src/logical_plan/plan.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index d91c38b828442..bda03fb7087a9 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -2131,7 +2131,7 @@ impl Filter { } Ok(Self { - predicate: predicate.unalias(), + predicate: predicate.unalias_nested().data, input, }) }