@@ -41,15 +41,15 @@ pub use order_by::rewrite_sort_cols_by_aggs;
4141
4242/// Trait for rewriting [`Expr`]s into function calls.
4343///
44- /// This trait is used with `FunctionRegistry::register_function_rewrite` to
44+ /// This trait is used with `FunctionRegistry::register_function_rewrite`
4545/// to evaluating `Expr`s using functions that may not be built in to DataFusion
4646///
4747/// For example, concatenating arrays `a || b` is represented as
4848/// `Operator::ArrowAt`, but can be implemented by calling a function
4949/// `array_concat` from the `functions-nested` crate.
5050// This is not used in datafusion internally, but it is still helpful for downstream project so don't remove it.
5151pub trait FunctionRewrite : Debug {
52- /// Return a human readable name for this rewrite
52+ /// Return a human- readable name for this rewrite
5353 fn name ( & self ) -> & str ;
5454
5555 /// Potentially rewrite `expr` to some other expression
@@ -219,26 +219,30 @@ pub fn strip_outer_reference(expr: Expr) -> Expr {
219219/// Returns plan with expressions coerced to types compatible with
220220/// schema types
221221pub fn coerce_plan_expr_for_schema (
222- plan : LogicalPlan ,
222+ plan : impl Into < Arc < LogicalPlan > > ,
223223 schema : & DFSchema ,
224- ) -> Result < LogicalPlan > {
225- match plan {
224+ ) -> Result < Arc < LogicalPlan > > {
225+ let plan = plan. into ( ) ;
226+ if matches ! ( plan. as_ref( ) , LogicalPlan :: Projection ( _) ) {
226227 // special case Projection to avoid adding multiple projections
227- LogicalPlan :: Projection ( Projection { expr, input, .. } ) => {
228- let new_exprs = coerce_exprs_for_schema ( expr, input. schema ( ) , schema) ?;
229- let projection = Projection :: try_new ( new_exprs, input) ?;
230- Ok ( LogicalPlan :: Projection ( projection) )
231- }
232- _ => {
233- let exprs: Vec < Expr > = plan. schema ( ) . iter ( ) . map ( Expr :: from) . collect ( ) ;
234- let new_exprs = coerce_exprs_for_schema ( exprs, plan. schema ( ) , schema) ?;
235- let add_project = new_exprs. iter ( ) . any ( |expr| expr. try_as_col ( ) . is_none ( ) ) ;
236- if add_project {
237- let projection = Projection :: try_new ( new_exprs, Arc :: new ( plan) ) ?;
238- Ok ( LogicalPlan :: Projection ( projection) )
239- } else {
240- Ok ( plan)
241- }
228+ let LogicalPlan :: Projection ( Projection { expr, input, .. } ) =
229+ Arc :: unwrap_or_clone ( plan)
230+ else {
231+ unreachable ! ( )
232+ } ;
233+
234+ let new_exprs = coerce_exprs_for_schema ( expr, input. schema ( ) , schema) ?;
235+ let projection = Projection :: try_new ( new_exprs, input) ?;
236+ Ok ( Arc :: new ( LogicalPlan :: Projection ( projection) ) )
237+ } else {
238+ let exprs: Vec < Expr > = plan. schema ( ) . iter ( ) . map ( Expr :: from) . collect ( ) ;
239+ let new_exprs = coerce_exprs_for_schema ( exprs, plan. schema ( ) , schema) ?;
240+ let add_project = new_exprs. iter ( ) . any ( |expr| expr. try_as_col ( ) . is_none ( ) ) ;
241+ if add_project {
242+ let projection = Projection :: try_new ( new_exprs, plan) ?;
243+ Ok ( Arc :: new ( LogicalPlan :: Projection ( projection) ) )
244+ } else {
245+ Ok ( plan)
242246 }
243247 }
244248}
@@ -442,7 +446,7 @@ mod test {
442446 fn normalize_cols ( ) {
443447 let expr = col ( "a" ) + col ( "b" ) + col ( "c" ) ;
444448
445- // Schemas with some matching and some non matching cols
449+ // Schemas with some matching and some non- matching cols
446450 let schema_a = make_schema_with_empty_metadata (
447451 vec ! [ Some ( "tableA" . into( ) ) , Some ( "tableA" . into( ) ) ] ,
448452 vec ! [ "a" , "aa" ] ,
0 commit comments