fix: move coercion of union from builder to `TypeCoercion` #11961

jonahgao · 2024-08-13T08:57:51Z

Which issue does this PR close?

Rationale for this change

Currently, in the LogicalPlan builder, we coerce the inputs of a union plan to a common schema.
But this common schema is inaccurate because it lacks type coercion for expressions.

For example, in #11742, the return type of the expression nvl(v1, 0.5) is INT before type coercion and will become Float64 afterward.

/// The return type of `nvl` depends on the type of the first argument.
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
    Ok(arg_types[0].clone())
}

The fix is to place coercion of the union after the type coercion of expressions.

What changes are included in this PR?

Move coercion of union from builder to TypeCoercion.
Move some union-related SQL planner tests to slt, as these tests need to perform type coercion to achieve the same results as before.

Are these changes tested?

Yes

Are there any user-facing changes?

Yes

project_with_column_index has become private. It was introduced for UNION in #2108, and I think it should only be used internally.

jonahgao · 2024-08-13T08:59:18Z

datafusion/expr/src/logical_plan/builder.rs

@@ -1881,23 +1799,6 @@ mod tests {
        Ok(())
    }

-    #[test]
-    fn plan_builder_union_different_num_columns_error() -> Result<()> {


Moved to slt.

jonahgao · 2024-08-13T09:01:01Z

datafusion/optimizer/src/eliminate_nested_union.rs

@@ -127,7 +130,14 @@ mod tests {
    }

    fn assert_optimized_plan_equal(plan: LogicalPlan, expected: &str) -> Result<()> {
-        assert_optimized_plan_eq(Arc::new(EliminateNestedUnion::new()), plan, expected)
+        let options = ConfigOptions::default();
+        let analyzed_plan = Analyzer::with_rules(vec![Arc::new(TypeCoercion::new())])


Add TypeCoercion to avoid breaking the tests.

jonahgao · 2024-08-13T09:04:34Z

datafusion/sqllogictest/test_files/type_coercion.slt

+05)--------Projection: Int64(1) AS a
+06)----------EmptyRelation
+07)--Projection: Float64(2.1) + x.a AS Float64(0) + x.a
+08)----Aggregate: groupBy=[[Float64(2.1) + CAST(x.a AS Float64)]], aggr=[[]]


This is the only difference I observed from the previous test results because type coercion was performed, and x.a had an additional cast.

jonahgao · 2024-08-13T09:16:00Z

datafusion/sql/tests/sql_integration.rs

@@ -2176,148 +2176,6 @@ fn union_all() {
    quick_test(sql, expected);
 }

-#[test]
-fn union_with_different_column_names() {


Converting these union tests to SLT is easier than moving them to TypeCoercion.

alamb

Thank you @jonahgao -- I agree it makes much more sense to have the type coercion for union happening as part of TypeCoercsion rather than the SQL planner / builder.

alamb · 2024-08-13T19:21:39Z

datafusion/expr/src/logical_plan/builder.rs

-        return plan_err!("Empty UNION");
-    }
-
+    // Temporarily use the schema from the left input and later rely on the analyzer to


alamb · 2024-08-13T19:27:02Z

datafusion/sqllogictest/test_files/type_coercion.slt

+05)----EmptyRelation
+
+# union_with_incompatible_data_type()
+query error Error during planning: UNION Column 'Int64\(1\)' \(type: Int64\) is not compatible with other type: Interval\(MonthDayNano\)


I noticed that the error refers to the types in reverse order than they appear in the query

The error message might be better if it were something more like

Incompatible inputs for Union. Previous inputs were of type Interval\(MonthDayNano\), got incomaptible type 'Int64\(1\)' \(type: Int64\)

Done. Thanks @alamb

alamb · 2024-08-14T11:17:48Z

datafusion/optimizer/src/analyzer/type_coercion.rs

@@ -809,7 +809,7 @@ fn coerce_case_expression(case: Case, schema: &DFSchema) -> Result<Case> {
 }

 /// Get a common schema that is compatible with all inputs of UNION.
-fn coerce_union_schema(inputs: Vec<Arc<LogicalPlan>>) -> Result<DFSchema> {
+fn coerce_union_schema(inputs: &Vec<Arc<LogicalPlan>>) -> Result<DFSchema> {


alamb · 2024-08-14T11:18:03Z

Looks great @jonahgao -- thank you

jonahgao · 2024-08-14T13:30:51Z

Thanks for the review @alamb

…pache#11961)" This reverts commit afa23ab.

jonahgao added 10 commits August 9, 2024 23:22

Improve type coercion for UNION

e019ccc

fix clippy

9fca957

fix test

e69a1a1

Merge branch 'main' into union_coercion

6c3edd3

fix sqllogictests

b67f11c

fix EliminateNestedUnion tests

5ad7243

Move tests to slt

d96892c

Merge branch 'main' into union_coercion

d672bbf

Move union_coercion to type_coercion.rs

6c05d17

fix tests

21f93dc

github-actions bot added sql SQL Planner logical-expr Logical plan and expressions optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) labels Aug 13, 2024

jonahgao commented Aug 13, 2024

View reviewed changes

fix cargo doc

c2235df

jonahgao commented Aug 13, 2024

View reviewed changes

alamb approved these changes Aug 13, 2024

View reviewed changes

jonahgao added 5 commits August 14, 2024 09:57

Merge branch 'main' into union_coercion

364ba12

Improve error msg

e5c2dc9

As static member

723fd43

Avoid clone

7e8b3f4

Fix clippy

d42486e

alamb approved these changes Aug 14, 2024

View reviewed changes

alamb merged commit afa23ab into apache:main Aug 14, 2024
24 checks passed

jonahgao deleted the union_coercion branch August 14, 2024 13:30

wiedld added a commit to influxdata/arrow-datafusion that referenced this pull request Aug 19, 2024

Revert "fix: move coercion of union from builder to TypeCoercion (a…

4885e32

…pache#11961)" This reverts commit afa23ab.

wiedld added a commit to influxdata/arrow-datafusion that referenced this pull request Aug 21, 2024

Revert "fix: move coercion of union from builder to TypeCoercion (a…

a4d9d28

…pache#11961)" This reverts commit afa23ab.

Omega359 mentioned this pull request Oct 24, 2024

Datafusion 42 does not raise plan error when UNION branches have different number of columns #13092

Closed

findepi mentioned this pull request Oct 26, 2024

Raise a plan error on union if column count is not the same between plans #13117

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: move coercion of union from builder to `TypeCoercion` #11961

fix: move coercion of union from builder to `TypeCoercion` #11961

jonahgao commented Aug 13, 2024

jonahgao Aug 13, 2024

jonahgao Aug 13, 2024

jonahgao Aug 13, 2024

jonahgao Aug 13, 2024

alamb left a comment

alamb Aug 13, 2024

alamb Aug 13, 2024

jonahgao Aug 14, 2024

alamb Aug 14, 2024

alamb commented Aug 14, 2024

jonahgao commented Aug 14, 2024

fix: move coercion of union from builder to TypeCoercion #11961

fix: move coercion of union from builder to TypeCoercion #11961

Conversation

jonahgao commented Aug 13, 2024

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb commented Aug 14, 2024

jonahgao commented Aug 14, 2024

fix: move coercion of union from builder to `TypeCoercion` #11961

fix: move coercion of union from builder to `TypeCoercion` #11961