Refactor Async UDF rewriting in physical planner

Is your feature request related to a problem or challenge?

#14837 introduced support for async UDFs. As part of it, in the physical planner it rewrites projections, filters (and inputs to aggregations by #17619).

Projection:

datafusion/datafusion/core/src/physical_planner.rs

Lines 2244 to 2271 in e323357

    
           match self.try_plan_async_exprs( 
        
               num_input_columns, 
        
               PlannedExprResult::ExprWithName(physical_exprs), 
        
               input_physical_schema.as_ref(), 
        
           )? { 
        
               PlanAsyncExpr::Sync(PlannedExprResult::ExprWithName(physical_exprs)) => { 
        
                   let proj_exprs: Vec<ProjectionExpr> = physical_exprs 
        
                       .into_iter() 
        
                       .map(|(expr, alias)| ProjectionExpr { expr, alias }) 
        
                       .collect(); 
        
                   Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input_exec)?)) 
        
               } 
        
               PlanAsyncExpr::Async( 
        
                   async_map, 
        
                   PlannedExprResult::ExprWithName(physical_exprs), 
        
               ) => { 
        
                   let async_exec = 
        
                       AsyncFuncExec::try_new(async_map.async_exprs, input_exec)?; 
        
                   let proj_exprs: Vec<ProjectionExpr> = physical_exprs 
        
                       .into_iter() 
        
                       .map(|(expr, alias)| ProjectionExpr { expr, alias }) 
        
                       .collect(); 
        
                   let new_proj_exec = 
        
                       ProjectionExec::try_new(proj_exprs, Arc::new(async_exec))?; 
        
                   Ok(Arc::new(new_proj_exec)) 
        
               } 
        
               _ => internal_err!("Unexpected PlanAsyncExpressions variant"), 
        
           }

Filter:

datafusion/datafusion/core/src/physical_planner.rs

Lines 834 to 875 in e323357

    
           LogicalPlan::Filter(Filter { 
        
               predicate, input, .. 
        
           }) => { 
        
               let physical_input = children.one()?; 
        
               let input_dfschema = input.schema(); 
        
               let runtime_expr = 
        
                   self.create_physical_expr(predicate, input_dfschema, session_state)?; 
        
               let input_schema = input.schema(); 
        
               let filter = match self.try_plan_async_exprs( 
        
                   input_schema.fields().len(), 
        
                   PlannedExprResult::Expr(vec![runtime_expr]), 
        
                   input_schema.as_arrow(), 
        
               )? { 
        
                   PlanAsyncExpr::Sync(PlannedExprResult::Expr(runtime_expr)) => { 
        
                       FilterExec::try_new(Arc::clone(&runtime_expr[0]), physical_input)? 
        
                   } 
        
                   PlanAsyncExpr::Async( 
        
                       async_map, 
        
                       PlannedExprResult::Expr(runtime_expr), 
        
                   ) => { 
        
                       let async_exec = AsyncFuncExec::try_new( 
        
                           async_map.async_exprs, 
        
                           physical_input, 
        
                       )?; 
        
                       FilterExec::try_new( 
        
                           Arc::clone(&runtime_expr[0]), 
        
                           Arc::new(async_exec), 
        
                       )? 
        
                       // project the output columns excluding the async functions 
        
                       // The async functions are always appended to the end of the schema. 
        
                       .with_projection(Some( 
        
                           (0..input.schema().fields().len()).collect(), 
        
                       ))? 
        
                   } 
        
                   _ => { 
        
                       return internal_err!( 
        
                           "Unexpected result from try_plan_async_exprs" 
        
                       ) 
        
                   } 
        
               };

Aggregation input:

datafusion/datafusion/core/src/physical_planner.rs

Lines 745 to 785 in e323357

    
           for agg_func in &mut aggregates { 
        
               match self.try_plan_async_exprs( 
        
                   num_input_columns, 
        
                   PlannedExprResult::Expr(agg_func.expressions()), 
        
                   physical_input_schema.as_ref(), 
        
               )? { 
        
                   PlanAsyncExpr::Async( 
        
                       async_map, 
        
                       PlannedExprResult::Expr(physical_exprs), 
        
                   ) => { 
        
                       async_exprs.extend(async_map.async_exprs); 
        
                       if let Some(new_agg_func) = agg_func.with_new_expressions( 
        
                           physical_exprs, 
        
                           agg_func 
        
                               .order_bys() 
        
                               .iter() 
        
                               .cloned() 
        
                               .map(|x| x.expr) 
        
                               .collect(), 
        
                       ) { 
        
                           *agg_func = Arc::new(new_agg_func); 
        
                       } else { 
        
                           return internal_err!("Failed to plan async expression"); 
        
                       } 
        
                   } 
        
                   PlanAsyncExpr::Sync(PlannedExprResult::Expr(_)) => { 
        
                       // Do nothing 
        
                   } 
        
                   _ => { 
        
                       return internal_err!( 
        
                           "Unexpected result from try_plan_async_exprs" 
        
                       ) 
        
                   } 
        
               } 
        
           } 
        
           let input_exec = if !async_exprs.is_empty() { 
        
               Arc::new(AsyncFuncExec::try_new(async_exprs, input_exec)?) 
        
           } else { 
        
               input_exec 
        
           };

I find this quite confusing to understand and it can cause issues (see #18149).

There was also discussion about this in the original async UDF function:

Describe the solution you'd like

Either we try to clean up this rewrite logic and do it earlier, or see if we can design in a way to not need to rewrite at all (see original PR as there seemed to be some discussion towards that direction).

Describe alternatives you've considered

No response

Additional context

No response

@Jefffrey

hey @Jefffrey if this Issue is open then assign to me i want to work on this...

@GaneshPatil7517

hey @Jefffrey if this Issue is open then assign to me i want to work on this...

@GaneshPatil7517 you can assign yourself to tickets you intend to tackle via a single take comment, see our contributor guide: https://datafusion.apache.org/contributor-guide/index.html#open-contribution-and-assigning-tickets

However I would caution you to properly investigate this to determine if it is achievable as it isn't exactly clear cut on what the solution is here.

take

@Jefffrey

@Jefffrey I was thinking of going in the direction of adding an extract_async_exprs() method to the ExecutionPlan trait, which would be similar pattern of gather_filters_for_pushdown(), where each operator implements its own async expression extraction. A simple physical optimizer rule (ExtractAsyncFuncs) would then call this method on each node and insert
AsyncFuncExec where needed.

This would be the logic per-operator, will alos fixe the #18149 schema mismatch bug (by properly reconstructing aggregates), and removes async handling complexity from the physical planner.

Did you give this any thought? I can explore further if you have any feedback on this.

@Jefffrey I was thinking of going in the direction of adding an extract_async_exprs() method to the ExecutionPlan trait, which would be similar pattern of gather_filters_for_pushdown(), where each operator implements its own async expression extraction. A simple physical optimizer rule (ExtractAsyncFuncs) would then call this method on each node and insert AsyncFuncExec where needed.

This would be the logic per-operator, will alos fixe the #18149 schema mismatch bug (by properly reconstructing aggregates), and removes async handling complexity from the physical planner.

Did you give this any thought? I can explore further if you have any feedback on this.

That does sound like a potential avenue to explore; my only concern is it would be a physical optimizer rule that would become essential. I don't know what the stance is on physical optimizer rules, if they should stick purely to optimizing or if its fine that they act as rewriting into valid plans 🤔

Thanks for raising this that makes sense. My goal was mainly to simplify async UDF handling by moving operator-specific logic out of the physical planner and fixing the schema mismatch issue (#18149). I understand the concern around optimizer rules becoming essential for plan correctness, and I’d really appreciate guidance on whether this direction aligns with DataFusion’s intended architecture or if there’s a better approach I should explore.........

@Jefffrey I was thinking of going in the direction of adding an extract_async_exprs() method to the ExecutionPlan trait, which would be similar pattern of gather_filters_for_pushdown(), where each operator implements its own async expression extraction. A simple physical optimizer rule (ExtractAsyncFuncs) would then call this method on each node and insert AsyncFuncExec where needed.
This would be the logic per-operator, will alos fixe the #18149 schema mismatch bug (by properly reconstructing aggregates), and removes async handling complexity from the physical planner.
Did you give this any thought? I can explore further if you have any feedback on this.

That does sound like a potential avenue to explore; my only concern is it would be a physical optimizer rule that would become essential. I don't know what the stance is on physical optimizer rules, if they should stick purely to optimizing or if its fine that they act as rewriting into valid plans 🤔

We can follow the precedent from the EnforceDistribution/EnforceSorting and add ExtractAsyncFuncs as an essential optimizer rule, but as you said it will add another essential rule. We can go in other direction and keep the async handling as is with just some refactoring, we can avoid the essential optimizer rules, I don't know how it improves the current situation apart from cleaner codes. If we want broader change we can make the planner create a marker wrapper around inputs with async expressions; then the optimizer will unwraps/transform them.

Would be possible to, during logical planning, move async udfs into a LogicalPlan::Projection below the original node, so the physical planner can only care about async udf in projections? This also would introduce the async columns (__async_fn_0) earlier during planning, which could fix #18149 and the second error I faced at #22091

Rewrite example from #18149

Aggregate: groupBy=[[]], aggr=[[approx_distinct(custom_udf(CAST(data.x AS Int64)))]]
    TableScan: data projection=[x]

Into:

Aggregate: groupBy=[[]], aggr=[[approx_distinct(__async_fn_0)]]
    Projection: custom_udf(CAST(data.x AS Int64)) AS __async_fn_0
        TableScan: data projection=[x]

Would be possible to, during logical planning, move async udfs into a LogicalPlan::Projection below the original node, so the physical planner can only care about async udf in projections? This also would introduce the async columns (__async_fn_0) earlier during planning, which could fix #18149 and the second error I faced at #22091

Rewrite example from #18149
Aggregate: groupBy=[[]], aggr=[[approx_distinct(custom_udf(CAST(data.x AS Int64)))]]
    TableScan: data projection=[x]
Into:
Aggregate: groupBy=[[]], aggr=[[approx_distinct(__async_fn_0)]]
    Projection: custom_udf(CAST(data.x AS Int64)) AS __async_fn_0
        TableScan: data projection=[x]

Would this potentially require modifying multiple optimizer rules to ensure they don't get pushed down?

Would this potentially require modifying multiple optimizer rules to ensure they don't get pushed down?

The async_udf being pulled again to the aggr node? Or the new projection node being pushed further down the tree? (Can this be a problem?) But anyway, you have a good point, ProjectionSimplifier would merge two consecutive projections (which is how the error I hit would be fixed), and that's only taking into account core rules, this could require changes in downstream rules, which is not great. This also wouldn't work for joins where the async_udf can use inputs from both sides.

... clean up this rewrite logic and do it earlier ...

I read the code again and I believe, as you said, that doing it earlier during physical planning can fix both issues, by creating the physical expressions using the schema that it's only available after creating the AsyncMapper, and maybe rewriting the logical expressions instead of the physical ones

Navigation Menu

Repository navigation

Refactor Async UDF rewriting in physical planner #18150

Is your feature request related to a problem or challenge?

Describe the solution you'd like

Describe alternatives you've considered

Additional context

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	match self.try_plan_async_exprs(
	num_input_columns,
	PlannedExprResult::ExprWithName(physical_exprs),
	input_physical_schema.as_ref(),
	)? {
	PlanAsyncExpr::Sync(PlannedExprResult::ExprWithName(physical_exprs)) => {
	let proj_exprs: Vec<ProjectionExpr> = physical_exprs
	.into_iter()
	.map(\|(expr, alias)\| ProjectionExpr { expr, alias })
	.collect();
	Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input_exec)?))
	}
	PlanAsyncExpr::Async(
	async_map,
	PlannedExprResult::ExprWithName(physical_exprs),
	) => {
	let async_exec =
	AsyncFuncExec::try_new(async_map.async_exprs, input_exec)?;
	let proj_exprs: Vec<ProjectionExpr> = physical_exprs
	.into_iter()
	.map(\|(expr, alias)\| ProjectionExpr { expr, alias })
	.collect();
	let new_proj_exec =
	ProjectionExec::try_new(proj_exprs, Arc::new(async_exec))?;
	Ok(Arc::new(new_proj_exec))
	}
	_ => internal_err!("Unexpected PlanAsyncExpressions variant"),
	}

	LogicalPlan::Filter(Filter {
	predicate, input, ..
	}) => {
	let physical_input = children.one()?;
	let input_dfschema = input.schema();

	let runtime_expr =
	self.create_physical_expr(predicate, input_dfschema, session_state)?;

	let input_schema = input.schema();
	let filter = match self.try_plan_async_exprs(
	input_schema.fields().len(),
	PlannedExprResult::Expr(vec![runtime_expr]),
	input_schema.as_arrow(),
	)? {
	PlanAsyncExpr::Sync(PlannedExprResult::Expr(runtime_expr)) => {
	FilterExec::try_new(Arc::clone(&runtime_expr[0]), physical_input)?
	}
	PlanAsyncExpr::Async(
	async_map,
	PlannedExprResult::Expr(runtime_expr),
	) => {
	let async_exec = AsyncFuncExec::try_new(
	async_map.async_exprs,
	physical_input,
	)?;
	FilterExec::try_new(
	Arc::clone(&runtime_expr[0]),
	Arc::new(async_exec),
	)?
	// project the output columns excluding the async functions
	// The async functions are always appended to the end of the schema.
	.with_projection(Some(
	(0..input.schema().fields().len()).collect(),
	))?
	}
	_ => {
	return internal_err!(
	"Unexpected result from try_plan_async_exprs"
	)
	}
	};

	for agg_func in &mut aggregates {
	match self.try_plan_async_exprs(
	num_input_columns,
	PlannedExprResult::Expr(agg_func.expressions()),
	physical_input_schema.as_ref(),
	)? {
	PlanAsyncExpr::Async(
	async_map,
	PlannedExprResult::Expr(physical_exprs),
	) => {
	async_exprs.extend(async_map.async_exprs);

	if let Some(new_agg_func) = agg_func.with_new_expressions(
	physical_exprs,
	agg_func
	.order_bys()
	.iter()
	.cloned()
	.map(\|x\| x.expr)
	.collect(),
	) {
	*agg_func = Arc::new(new_agg_func);
	} else {
	return internal_err!("Failed to plan async expression");
	}
	}
	PlanAsyncExpr::Sync(PlannedExprResult::Expr(_)) => {
	// Do nothing
	}
	_ => {
	return internal_err!(
	"Unexpected result from try_plan_async_exprs"
	)
	}
	}
	}
	let input_exec = if !async_exprs.is_empty() {
	Arc::new(AsyncFuncExec::try_new(async_exprs, input_exec)?)
	} else {
	input_exec
	};

Refactor Async UDF rewriting in physical planner #18150

Description

Is your feature request related to a problem or challenge?

Describe the solution you'd like

Describe alternatives you've considered

Additional context

Activity

GaneshPatil7517 commented on Jan 12, 2026

Jefffrey commented on Jan 12, 2026

GaneshPatil7517 commented on Jan 13, 2026

kumarUjjawal commented on Jan 16, 2026

Jefffrey commented on Jan 16, 2026

GaneshPatil7517 commented on Jan 16, 2026

kumarUjjawal commented on Jan 17, 2026

gstvg commented on May 10, 2026

Jefffrey commented on May 11, 2026

gstvg commented on May 15, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions