@@ -102,15 +102,17 @@ fn analyze_internal(
102
102
// select t2.c2 from t1 where t1.c1 in (select t2.c1 from t2 where t2.c2=t1.c3)
103
103
schema. merge ( external_schema) ;
104
104
105
- let mut expr_rewrite = TypeCoercionRewriter { schema : & schema } ;
105
+ let mut expr_rewrite = TypeCoercionRewriter :: new ( & schema) ;
106
106
107
107
let name_preserver = NamePreserver :: new ( & plan) ;
108
- // apply coercion rewrite all expressions in the plan
108
+ // apply coercion rewrite all expressions in the plan indivdually
109
109
plan. map_expressions ( |expr| {
110
110
let original_name = name_preserver. save ( & expr) ?;
111
111
expr. rewrite ( & mut expr_rewrite) ?
112
112
. map_data ( |expr| original_name. restore ( expr) )
113
113
} ) ?
114
+ // coerce join expressions specially
115
+ . map_data ( |plan| expr_rewrite. coerce_joins ( plan) ) ?
114
116
// recompute the schema after the expressions have been rewritten as the types may have changed
115
117
. map_data ( |plan| plan. recompute_schema ( ) )
116
118
}
@@ -119,6 +121,56 @@ pub(crate) struct TypeCoercionRewriter<'a> {
119
121
pub ( crate ) schema : & ' a DFSchema ,
120
122
}
121
123
124
+ impl < ' a > TypeCoercionRewriter < ' a > {
125
+ fn new ( schema : & ' a DFSchema ) -> Self {
126
+ Self { schema }
127
+ }
128
+
129
+ /// Coerce join equality expressions
130
+ ///
131
+ /// Joins must be treated specially as their equality expressions are stored
132
+ /// as a parallel list of left and right expressions, rather than a single
133
+ /// equality expression
134
+ ///
135
+ /// For example, on_exprs like `t1.a = t2.b AND t1.x = t2.y` will be stored
136
+ /// as a list of `(t1.a, t2.b), (t1.x, t2.y)`
137
+ fn coerce_joins ( & mut self , plan : LogicalPlan ) -> Result < LogicalPlan > {
138
+ let LogicalPlan :: Join ( mut join) = plan else {
139
+ return Ok ( plan) ;
140
+ } ;
141
+
142
+ join. on = join
143
+ . on
144
+ . into_iter ( )
145
+ . map ( |( lhs, rhs) | {
146
+ // coerce the arguments as though they were a single binary equality
147
+ // expression
148
+ let ( lhs, rhs) = self . coerce_binary_op ( lhs, Operator :: Eq , rhs) ?;
149
+ Ok ( ( lhs, rhs) )
150
+ } )
151
+ . collect :: < Result < Vec < _ > > > ( ) ?;
152
+
153
+ Ok ( LogicalPlan :: Join ( join) )
154
+ }
155
+
156
+ fn coerce_binary_op (
157
+ & self ,
158
+ left : Expr ,
159
+ op : Operator ,
160
+ right : Expr ,
161
+ ) -> Result < ( Expr , Expr ) > {
162
+ let ( left_type, right_type) = get_input_types (
163
+ & left. get_type ( self . schema ) ?,
164
+ & op,
165
+ & right. get_type ( self . schema ) ?,
166
+ ) ?;
167
+ Ok ( (
168
+ left. cast_to ( & left_type, self . schema ) ?,
169
+ right. cast_to ( & right_type, self . schema ) ?,
170
+ ) )
171
+ }
172
+ }
173
+
122
174
impl < ' a > TreeNodeRewriter for TypeCoercionRewriter < ' a > {
123
175
type Node = Expr ;
124
176
@@ -223,15 +275,11 @@ impl<'a> TreeNodeRewriter for TypeCoercionRewriter<'a> {
223
275
) ) ) )
224
276
}
225
277
Expr :: BinaryExpr ( BinaryExpr { left, op, right } ) => {
226
- let ( left_type, right_type) = get_input_types (
227
- & left. get_type ( self . schema ) ?,
228
- & op,
229
- & right. get_type ( self . schema ) ?,
230
- ) ?;
278
+ let ( left, right) = self . coerce_binary_op ( * left, op, * right) ?;
231
279
Ok ( Transformed :: yes ( Expr :: BinaryExpr ( BinaryExpr :: new (
232
- Box :: new ( left. cast_to ( & left_type , self . schema ) ? ) ,
280
+ Box :: new ( left) ,
233
281
op,
234
- Box :: new ( right. cast_to ( & right_type , self . schema ) ? ) ,
282
+ Box :: new ( right) ,
235
283
) ) ) )
236
284
}
237
285
Expr :: Between ( Between {
0 commit comments