alamb commented on code in PR #3513:
URL: https://github.com/apache/arrow-datafusion/pull/3513#discussion_r973314487
##########
datafusion/sql/src/planner.rs:
##########
@@ -4171,6 +4171,39 @@ mod tests {
);
}
+ #[test]
+ fn union_with_different_decimal_data_types() {
+ let sql = "SELECT 1 a UNION ALL SELECT 1.1 a";
+ let expected = "Union\
+ \n Projection: CAST(Int64(1) AS a AS Float64) AS a\
+ \n EmptyRelation\
+ \n Projection: Float64(1.1) AS a\
+ \n EmptyRelation";
+ quick_test(sql, expected);
+ }
+
+ #[test]
+ fn union_with_null() {
+ let sql = "SELECT NULL a UNION ALL SELECT 1.1 a";
+ let expected = "Union\
+ \n Projection: CAST(NULL AS a AS Float64) AS a\
+ \n EmptyRelation\
+ \n Projection: Float64(1.1) AS a\
+ \n EmptyRelation";
+ quick_test(sql, expected);
+ }
+
+ #[test]
+ fn union_with_float_and_string() {
+ let sql = "SELECT 'a' a UNION ALL SELECT 1.1 a";
+ let expected = "Union\
+ \n Projection: Utf8(\"a\") AS a\
+ \n EmptyRelation\
+ \n Projection: CAST(Float64(1.1) AS a AS Utf8) AS a\
+ \n EmptyRelation";
+ quick_test(sql, expected);
+ }
+
Review Comment:
Can you please add at least two more test cases here:
1. A test when the union has multiple columns (not just one)
2. A test when the union has uncoercable types (aka test the error path).
##########
datafusion/expr/src/logical_plan/builder.rs:
##########
@@ -882,43 +885,59 @@ pub fn union_with_alias(
right_plan: LogicalPlan,
alias: Option<String>,
) -> Result<LogicalPlan> {
- let union_schema = left_plan.schema().clone();
- let inputs_iter = vec![left_plan, right_plan]
+ let union_schema = (0..left_plan.schema().fields().len())
+ .map(|i| {
+ let left_field = left_plan.schema().field(i);
+ let right_field = right_plan.schema().field(i);
+ let nullable = left_field.is_nullable() ||
right_field.is_nullable();
+ let data_type =
+ comparison_coercion(left_field.data_type(),
right_field.data_type())
+ .ok_or_else(|| {
+ return DataFusionError::Plan(format!(
+ "Column {} (type: {}) is not compatible with column {}
(type: {})",
Review Comment:
```suggestion
"UNION Column {} (type: {}) is not compatible with
column {} (type: {})",
```
##########
datafusion/expr/src/expr_rewriter.rs:
##########
@@ -524,6 +524,31 @@ pub fn unnormalize_cols(exprs: impl IntoIterator<Item =
Expr>) -> Vec<Expr> {
exprs.into_iter().map(unnormalize_col).collect()
}
+/// Returns plan with expressions coerced to types compatible with
+/// schema types
+pub fn coerce_plan_expr_for_schema(
+ plan: &LogicalPlan,
+ schema: &DFSchema,
+) -> Result<LogicalPlan> {
+ let new_expr = plan
Review Comment:
This will work for plans that have a `LogicalPlan::Projection` , but I
wonder if it would work for other types of plans (like `LogicalExpr::GroupBy`)
where the expressions may be more constrained
Given this code is currently called just for union (which probably are
always `LogicalPlan::Projection`) I suspect this is probably fine.
Another potential implementation, is to add a new `LogicalPlan::Projection`
node (rather than rewriting the existing one) that does the casting.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
