alamb commented on code in PR #10356:
URL: https://github.com/apache/datafusion/pull/10356#discussion_r1598690159
##########
datafusion/optimizer/src/analyzer/type_coercion.rs:
##########
@@ -98,25 +101,75 @@ fn analyze_internal(
// select t2.c2 from t1 where t1.c1 in (select t2.c1 from t2 where
t2.c2=t1.c3)
schema.merge(external_schema);
- let mut expr_rewrite = TypeCoercionRewriter { schema: &schema };
-
- let new_expr = plan
- .expressions()
- .into_iter()
- .map(|expr| {
- // ensure aggregate names don't change:
- // https://github.com/apache/datafusion/issues/3555
- rewrite_preserving_name(expr, &mut expr_rewrite)
- })
- .collect::<Result<Vec<_>>>()?;
-
- plan.with_new_exprs(new_expr, new_inputs)
+ let mut expr_rewrite = TypeCoercionRewriter::new(&schema);
+
+ let name_preserver = NamePreserver::new(&plan);
+ // apply coercion rewrite all expressions in the plan indivdually
+ plan.map_expressions(|expr| {
+ let original_name = name_preserver.save(&expr)?;
+ expr.rewrite(&mut expr_rewrite)?
+ .map_data(|expr| original_name.restore(expr))
+ })?
+ // coerce join expressions specially
+ .map_data(|plan| expr_rewrite.coerce_joins(plan))?
+ // recompute the schema after the expressions have been rewritten as the
types may have changed
+ .map_data(|plan| plan.recompute_schema())
}
pub(crate) struct TypeCoercionRewriter<'a> {
pub(crate) schema: &'a DFSchema,
}
+impl<'a> TypeCoercionRewriter<'a> {
+ fn new(schema: &'a DFSchema) -> Self {
+ Self { schema }
+ }
+
+ /// Coerce join equality expressions
+ ///
+ /// Joins must be treated specially as their equality expressions are
stored
+ /// as a parallel list of left and right expressions, rather than a single
+ /// equality expression
+ ///
+ /// For example, on_exprs like `t1.a = t2.b AND t1.x = t2.y` will be stored
+ /// as a list of `(t1.a, t2.b), (t1.x, t2.y)`
+ fn coerce_joins(&mut self, plan: LogicalPlan) -> Result<LogicalPlan> {
+ let LogicalPlan::Join(mut join) = plan else {
+ return Ok(plan);
+ };
+
+ join.on = join
+ .on
+ .into_iter()
+ .map(|(lhs, rhs)| {
+ // coerce the arguments as though they were a single binary
equality
+ // expression
+ let (lhs, rhs) = self.coerce_binary_op(lhs, Operator::Eq,
rhs)?;
Review Comment:
I think `coerce_binary_op` is different than just casting `lhs` and `rhs`
as it first calls
[`get_input_types`](https://github.com/apache/datafusion/blob/9cc981b06115ee40b53384c287689ce0e07950bc/datafusion/expr/src/type_coercion/binary.rs#L204):
```rust
let (left_type, right_type) = get_input_types(
&left.get_type(self.schema)?,
&op,
&right.get_type(self.schema)?,
)?;
```
And `get_input_types` usese the comparison coercion rules to figure out a
common set if types to coerce `lhs` and `rhs` to.
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
