@@ -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
@@ -259,26 +259,30 @@ pub fn strip_outer_reference(expr: Expr) -> Expr {
259259/// Returns plan with expressions coerced to types compatible with
260260/// schema types
261261pub fn coerce_plan_expr_for_schema (
262- plan : LogicalPlan ,
262+ plan : impl Into < Arc < LogicalPlan > > ,
263263 schema : & DFSchema ,
264- ) -> Result < LogicalPlan > {
265- match plan {
264+ ) -> Result < Arc < LogicalPlan > > {
265+ let plan = plan. into ( ) ;
266+ if matches ! ( plan. as_ref( ) , LogicalPlan :: Projection ( _) ) {
266267 // special case Projection to avoid adding multiple projections
267- LogicalPlan :: Projection ( Projection { expr, input, .. } ) => {
268- let new_exprs = coerce_exprs_for_schema ( expr, input. schema ( ) , schema) ?;
269- let projection = Projection :: try_new ( new_exprs, input) ?;
270- Ok ( LogicalPlan :: Projection ( projection) )
271- }
272- _ => {
273- let exprs: Vec < Expr > = plan. schema ( ) . iter ( ) . map ( Expr :: from) . collect ( ) ;
274- let new_exprs = coerce_exprs_for_schema ( exprs, plan. schema ( ) , schema) ?;
275- let add_project = new_exprs. iter ( ) . any ( |expr| expr. try_as_col ( ) . is_none ( ) ) ;
276- if add_project {
277- let projection = Projection :: try_new ( new_exprs, Arc :: new ( plan) ) ?;
278- Ok ( LogicalPlan :: Projection ( projection) )
279- } else {
280- Ok ( plan)
281- }
268+ let LogicalPlan :: Projection ( Projection { expr, input, .. } ) =
269+ Arc :: unwrap_or_clone ( plan)
270+ else {
271+ unreachable ! ( )
272+ } ;
273+
274+ let new_exprs = coerce_exprs_for_schema ( expr, input. schema ( ) , schema) ?;
275+ let projection = Projection :: try_new ( new_exprs, input) ?;
276+ Ok ( Arc :: new ( LogicalPlan :: Projection ( projection) ) )
277+ } else {
278+ let exprs: Vec < Expr > = plan. schema ( ) . iter ( ) . map ( Expr :: from) . collect ( ) ;
279+ let new_exprs = coerce_exprs_for_schema ( exprs, plan. schema ( ) , schema) ?;
280+ let add_project = new_exprs. iter ( ) . any ( |expr| expr. try_as_col ( ) . is_none ( ) ) ;
281+ if add_project {
282+ let projection = Projection :: try_new ( new_exprs, plan) ?;
283+ Ok ( Arc :: new ( LogicalPlan :: Projection ( projection) ) )
284+ } else {
285+ Ok ( plan)
282286 }
283287 }
284288}
@@ -482,7 +486,7 @@ mod test {
482486 fn normalize_cols ( ) {
483487 let expr = col ( "a" ) + col ( "b" ) + col ( "c" ) ;
484488
485- // Schemas with some matching and some non matching cols
489+ // Schemas with some matching and some non- matching cols
486490 let schema_a = make_schema_with_empty_metadata (
487491 vec ! [ Some ( "tableA" . into( ) ) , Some ( "tableA" . into( ) ) ] ,
488492 vec ! [ "a" , "aa" ] ,
0 commit comments