Fix UNION field nullability tracking #14356
Conversation
github-actions
bot
added
logical-expr
This commit fixes two bugs related to UNION handling - when constructing union plan nullability of the other union branch was ignored, thus resulting field could easily have incorrect nullability - when pruning/simplifying projects, in `recompute_schema` function there was similar logic, thus loosing nullability information even for correctly constructed Union plan node As a result, other optimizer logic (e.g. `expr_simplifier.rs`) could draw incorrect conclusions and thus lead to incorrect query results, as demonstrated with the attached SLT test.
findepi
force-pushed
the
findepi/union
branch
from
b3c6f07 to
ada016c
Compare
|
cc @felipecrv, @Omega359 |
|
@wiedld can you also please review this PR? |
| /// Constructs new Union instance deriving schema from inputs. | ||
| /// | ||
| /// `loose_types` if true, inputs do not have to have matching types and produced schema will | ||
| /// take type from the first input. TODO this is not necessarily reasonable behavior. |
There was a problem hiding this comment.
Possible alternative: coerce type to common type? https://www.postgresql.org/docs/17/typeconv-union-case.html is what PG does
There was a problem hiding this comment.
I agree this is what should be happening.
There are two places where this code is run:
- plan building
- recompute_schema
For now in the (1) case I retained pre-existing logic, see current main
datafusion/datafusion/expr/src/logical_plan/builder.rs
Lines 1530 to 1531 in e718c1a
For (2) we don't want any coercions at all.
|
I ran the sqlite tests against this branch with no changes so the tests in there did not cover this case. |
|
Thanks @Omega359 for your review, addressed! |
| /// Constructs new Union instance deriving schema from inputs. | ||
| /// Inputs do not have to have matching types and produced schema will | ||
| /// take type from the first input. | ||
| pub fn try_new_with_loose_types(inputs: Vec<Arc<LogicalPlan>>) -> Result<Self> { |
There was a problem hiding this comment.
Another name for this one is try_new_with_coerce_types to emphasize it is coercing the types
There was a problem hiding this comment.
I believe this one should be coercing types, but it's not doing this. It just takes a type from first UNION branch.
This is not new logic though. It's just moved from plan builder over here, to avoid duplicating code.
I didn't know why this logic existed in this shape, but it felt intentional enough not to simply replace it in this PR. I would prefer it to be fixed later... Hence this function name to make the caller wonder what "loose types" mean.
There was a problem hiding this comment.
Yeah, this was noted once before. Thank you for making it more explicit.
| inputs: &[Arc<LogicalPlan>], | ||
| loose_types: bool, | ||
| ) -> Result<DFSchemaRef> { | ||
| if inputs.len() < 2 { |
There was a problem hiding this comment.
I think there is already code that computes the coerced schema in the analyzer:
Can we reuse the same logic? Maybe we can move the coercion code here
There was a problem hiding this comment.
In the optimizer case, no coercion logic should be invoked, the types must match.
So we need a version of this code which does that. (#14296 (comment))
For the variant which constructs new schema from some uncoerced inputs -- currently called "loose types" and currently not doing coercion to maintain behavior as it was -- yes, i agree this one could use the coerce_union_schema from the analyzer (perhaps moving the logic into here)
There was a problem hiding this comment.
maybe we can file a ticket describing what is desired and leave a comment in the code referring to that ticket
There was a problem hiding this comment.
Making sure I understand the use case. If I want to construct a UNION logical plan with different types that are coercible (be it by current builtin rules or future user-defined rules), then I would use the Union::try_new_with_loose_types and have the analyzer pass handle coercion. Is this right?
Then what exactly is the use case for the Union::try_new? Since it's used in the schema recompute which can occur after the analyzer type coercion. Do we therefore need it to always perform proper coercion, including any future user-defined coercion rules?
| inputs: &[Arc<LogicalPlan>], | ||
| loose_types: bool, | ||
| ) -> Result<DFSchemaRef> { | ||
| if inputs.len() < 2 { |
There was a problem hiding this comment.
maybe we can file a ticket describing what is desired and leave a comment in the code referring to that ticket
This commit fixes two bugs related to UNION handling
ignored, thus resulting field could easily have incorrect nullability
recompute_schemafunctionthere was similar logic, thus loosing nullability information even for
correctly constructed Union plan node
As a result, other optimizer logic (e.g.
expr_simplifier.rs) coulddraw incorrect conclusions and thus lead to incorrect query results, as
demonstrated with the attached SLT test.