From d30434a2a268ad7fdd9ee61008a51a3f515b188a Mon Sep 17 00:00:00 2001 From: xudong963 Date: Wed, 15 Jan 2025 09:07:30 +0800 Subject: [PATCH 1/4] Return err if wildcard is not expanded before type coercion --- .../optimizer/src/analyzer/type_coercion.rs | 25 +++++++++++-------- .../optimizer/tests/optimizer_integration.rs | 23 +++++++++++++++++ 2 files changed, 38 insertions(+), 10 deletions(-) diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index fe65af0a1486..725fab3a0df8 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -995,16 +995,19 @@ fn project_with_column_index( .enumerate() .map(|(i, e)| match e { Expr::Alias(Alias { ref name, .. }) if name != schema.field(i).name() => { - e.unalias().alias(schema.field(i).name()) + Ok(e.unalias().alias(schema.field(i).name())) } Expr::Column(Column { relation: _, ref name, - }) if name != schema.field(i).name() => e.alias(schema.field(i).name()), - Expr::Alias { .. } | Expr::Column { .. } => e, - _ => e.alias(schema.field(i).name()), + }) if name != schema.field(i).name() => Ok(e.alias(schema.field(i).name())), + Expr::Alias { .. } | Expr::Column { .. } => Ok(e), + Expr::Wildcard { .. } => { + plan_err!("Wildcard should be expanded before type coercion") + } + _ => Ok(e.alias(schema.field(i).name())), }) - .collect::>(); + .collect::>>()?; Projection::try_new_with_schema(alias_expr, input, schema) .map(LogicalPlan::Projection) @@ -1018,11 +1021,17 @@ mod test { use arrow::datatypes::DataType::Utf8; use arrow::datatypes::{DataType, Field, TimeUnit}; + use crate::analyzer::type_coercion::{ + coerce_case_expression, TypeCoercion, TypeCoercionRewriter, + }; + use crate::test::{assert_analyzed_plan_eq, assert_analyzed_plan_with_config_eq}; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{TransformedResult, TreeNode}; use datafusion_common::{DFSchema, DFSchemaRef, Result, ScalarValue}; use datafusion_expr::expr::{self, InSubquery, Like, ScalarFunction}; use datafusion_expr::logical_plan::{EmptyRelation, Projection, Sort}; + use datafusion_expr::sqlparser::dialect::PostgreSqlDialect; + use datafusion_expr::sqlparser::parser::Parser; use datafusion_expr::test::function_stub::avg_udaf; use datafusion_expr::{ cast, col, create_udaf, is_true, lit, AccumulatorFactoryFunction, AggregateUDF, @@ -1031,11 +1040,7 @@ mod test { Volatility, }; use datafusion_functions_aggregate::average::AvgAccumulator; - - use crate::analyzer::type_coercion::{ - coerce_case_expression, TypeCoercion, TypeCoercionRewriter, - }; - use crate::test::{assert_analyzed_plan_eq, assert_analyzed_plan_with_config_eq}; + use datafusion_sql::planner::SqlToRel; fn empty() -> Arc { Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { diff --git a/datafusion/optimizer/tests/optimizer_integration.rs b/datafusion/optimizer/tests/optimizer_integration.rs index 29fac5cc3dec..d1c3e84f83d0 100644 --- a/datafusion/optimizer/tests/optimizer_integration.rs +++ b/datafusion/optimizer/tests/optimizer_integration.rs @@ -23,10 +23,12 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use datafusion_common::config::ConfigOptions; use datafusion_common::{plan_err, Result}; +use datafusion_expr::sqlparser::dialect::PostgreSqlDialect; use datafusion_expr::test::function_stub::sum_udaf; use datafusion_expr::{AggregateUDF, LogicalPlan, ScalarUDF, TableSource, WindowUDF}; use datafusion_functions_aggregate::average::avg_udaf; use datafusion_functions_aggregate::count::count_udaf; +use datafusion_optimizer::analyzer::type_coercion::TypeCoercionRewriter; use datafusion_optimizer::analyzer::Analyzer; use datafusion_optimizer::optimizer::Optimizer; use datafusion_optimizer::{OptimizerConfig, OptimizerContext, OptimizerRule}; @@ -387,6 +389,27 @@ fn select_correlated_predicate_subquery_with_uppercase_ident() { assert_eq!(expected, format!("{plan}")); } +// The test should return an error +// because the wildcard didn't be expanded before type coercion +#[test] +fn test_union_coercion_with_wildcard() -> Result<()> { + let dialect = PostgreSqlDialect {}; + let context_provider = MyContextProvider::default(); + let sql = "select * from (SELECT col_int32, col_uint32 FROM test) union all select * from(SELECT col_uint32, col_int32 FROM test)"; + let statements = Parser::parse_sql(&dialect, sql)?; + let sql_to_rel = SqlToRel::new(&context_provider); + let logical_plan = sql_to_rel.sql_statement_to_plan(statements[0].clone())?; + + if let LogicalPlan::Union(union) = logical_plan { + let err = TypeCoercionRewriter::coerce_union(union) + .err() + .unwrap() + .to_string(); + assert_eq!(err, "Error during planning: Wildcard should be expanded before going to the method"); + } + Ok(()) +} + fn test_sql(sql: &str) -> Result { // parse the SQL let dialect = GenericDialect {}; // or AnsiDialect, or your own dialect ... From 791aafbc4ba9026c565ea8acd2589d999fef887d Mon Sep 17 00:00:00 2001 From: xudong963 Date: Wed, 15 Jan 2025 21:35:22 +0800 Subject: [PATCH 2/4] fix test --- datafusion/optimizer/tests/optimizer_integration.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/datafusion/optimizer/tests/optimizer_integration.rs b/datafusion/optimizer/tests/optimizer_integration.rs index d1c3e84f83d0..745c2354a025 100644 --- a/datafusion/optimizer/tests/optimizer_integration.rs +++ b/datafusion/optimizer/tests/optimizer_integration.rs @@ -22,7 +22,7 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use datafusion_common::config::ConfigOptions; -use datafusion_common::{plan_err, Result}; +use datafusion_common::{assert_contains, plan_err, Result}; use datafusion_expr::sqlparser::dialect::PostgreSqlDialect; use datafusion_expr::test::function_stub::sum_udaf; use datafusion_expr::{AggregateUDF, LogicalPlan, ScalarUDF, TableSource, WindowUDF}; @@ -405,7 +405,10 @@ fn test_union_coercion_with_wildcard() -> Result<()> { .err() .unwrap() .to_string(); - assert_eq!(err, "Error during planning: Wildcard should be expanded before going to the method"); + assert_contains!( + err, + "Error during planning: Wildcard should be expanded before type coercion" + ); } Ok(()) } From d17df1d8ec2fff401b3c054aa9bc8fb56289199b Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 15 Jan 2025 12:12:24 -0500 Subject: [PATCH 3/4] fix clippy --- datafusion/optimizer/src/analyzer/type_coercion.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 725fab3a0df8..48a5e2f9a07c 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -1030,8 +1030,6 @@ mod test { use datafusion_common::{DFSchema, DFSchemaRef, Result, ScalarValue}; use datafusion_expr::expr::{self, InSubquery, Like, ScalarFunction}; use datafusion_expr::logical_plan::{EmptyRelation, Projection, Sort}; - use datafusion_expr::sqlparser::dialect::PostgreSqlDialect; - use datafusion_expr::sqlparser::parser::Parser; use datafusion_expr::test::function_stub::avg_udaf; use datafusion_expr::{ cast, col, create_udaf, is_true, lit, AccumulatorFactoryFunction, AggregateUDF, @@ -1040,7 +1038,6 @@ mod test { Volatility, }; use datafusion_functions_aggregate::average::AvgAccumulator; - use datafusion_sql::planner::SqlToRel; fn empty() -> Arc { Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { From 48c7c51e141845d75aa00a1d90ef0fa3448994e9 Mon Sep 17 00:00:00 2001 From: xudong963 Date: Thu, 16 Jan 2025 07:27:27 +0800 Subject: [PATCH 4/4] improve test --- datafusion/optimizer/tests/optimizer_integration.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/datafusion/optimizer/tests/optimizer_integration.rs b/datafusion/optimizer/tests/optimizer_integration.rs index 745c2354a025..b9073f5ac881 100644 --- a/datafusion/optimizer/tests/optimizer_integration.rs +++ b/datafusion/optimizer/tests/optimizer_integration.rs @@ -409,6 +409,8 @@ fn test_union_coercion_with_wildcard() -> Result<()> { err, "Error during planning: Wildcard should be expanded before type coercion" ); + } else { + panic!("Expected Union plan"); } Ok(()) }