From e89d4c8a4836e1a68e4e4623083f030db7d5f2c1 Mon Sep 17 00:00:00 2001 From: Samuel Colvin Date: Mon, 8 Jul 2024 20:42:05 +0100 Subject: [PATCH] Remove any aliases in `Filter::try_new` rather than erroring (#11307) * allow alias in predicate, fix #11306 * Fix typo. Co-authored-by: Andrew Lamb * unalise predicate * use unalias_nested --------- Co-authored-by: Andrew Lamb --- .../user_defined/user_defined_sql_planner.rs | 37 ++++++++++++++++++- datafusion/expr/src/logical_plan/plan.rs | 16 +++----- 2 files changed, 41 insertions(+), 12 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..6b660e15a1970 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_select() { + 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..bda03fb7087a9 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,16 +2130,10 @@ 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 }) + Ok(Self { + predicate: predicate.unalias_nested().data, + input, + }) } /// Is this filter guaranteed to return 0 or 1 row in a given instantiation?