Repository navigation
Nesting async UDF calls causes an internal error #20031
Description
Describe the bug
Directly calling an async UDF on the output of another async UDF async_example(async_example(1)) produces an internal error.:
Internal error: async functions should not be called directly
They do not have to be the same async UDF in order to get this error.
To Reproduce
The following code is a miminal example somewhat copied from the async UDF example in this repo:
use std::sync::Arc;
use async_trait::async_trait;
use datafusion::{
arrow::{
array::{ArrayRef, Int64Array},
datatypes::DataType,
},
common::{cast::as_int64_array, not_impl_err, utils::take_function_args},
error::Result,
logical_expr::{
ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature,
async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl},
},
prelude::SessionContext,
};
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct AsyncExample {
signature: Signature,
}
impl ScalarUDFImpl for AsyncExample {
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn name(&self) -> &str {
"async_example"
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(DataType::Int64)
}
fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
not_impl_err!("async_udf was called without async")
}
}
#[async_trait]
impl AsyncScalarUDFImpl for AsyncExample {
async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
let args = ColumnarValue::values_to_arrays(&args.args)?;
let [input_column] = take_function_args(self.name(), args)?;
let input_column = as_int64_array(&input_column)?;
let results: Int64Array = input_column.iter().map(|i| i.map(|i| i + 1)).collect();
Ok(ColumnarValue::Array(Arc::new(results) as ArrayRef))
}
}
#[tokio::main]
async fn main() {
let ctx = SessionContext::new();
ctx.register_udf(
AsyncScalarUDF::new(Arc::new(AsyncExample {
signature: Signature::exact(
vec![DataType::Int64],
datafusion::logical_expr::Volatility::Immutable,
),
}))
.into_scalar_udf(),
);
// Works
let df = ctx.sql("SELECT async_example(1)").await.unwrap();
df.collect().await.unwrap();
// Does not work
let df = ctx
.sql("SELECT async_example(async_example(1))")
.await
.unwrap();
if let Err(err) = df.collect().await {
eprintln!("{err}");
}
}A backtrace for this example can be found below.
Expected behavior
Both UDFs should simply be evaluated in succession. if it this not yet supported a more specific error message may be useful.
Additional context
The above code produces the following output: (with the backtrace feature in this crate enabled)
I'm using datafusion 52.1.0
Activity
Tushar7012 commentedon Jan 27, 2026
I'm happy to work on implementing the fix if this analysis looks correct! Could I please be assigned to this issue? I'd like to proceed with the following plan as the immediate fix and add comprehensive test coverage for nested async UDF scenarios.
My proposed implementation plan:
1.Modify AsyncMapper::find_references() to stop recursion after finding an async UDF
2.Add defensive error handling in AsyncFuncExpr::invoke_with_args()
3.Create test cases for various nesting scenarios in
async_udf.slt
4.Update error messages to be more descriptive
Please let me know if you'd prefer a different approach or if there are any concerns with this solution!
Jefffrey commentedon Jan 28, 2026
These sound reasonable
Not sure what this means
Tushar7012 commentedon Feb 1, 2026
Hi @Jefffrey ,
By “defensive error handling in AsyncFuncExpr::invoke_with_args()”, I’m not suggesting changing the core execution logic or adding broad try/catch-style behavior. What I had in mind is a small, targeted guard to handle cases where a nested async UDF invocation reaches an unexpected or unsupported state (which is what currently results in the internal error).
Concretely, this would mean:
Detecting when invoke_with_args() is reached in an invalid nested async context
Failing fast with a clear, user-facing error message instead of hitting an internal error / panic
Keeping this logic minimal and only scoped to async UDF nesting, not general execution paths
That said, if you feel this guard is better placed elsewhere (or not needed once AsyncMapper::find_references() is fixed), I’m happy to drop or adjust this part and keep the change focused on the mapper + tests.
Jefffrey commentedon Feb 1, 2026
I don't quite understand what is being proposed; if we reach invoke_with_args then it already returns an internal error for us, so I don't know what more we need? This was indeed a bug and the internal error is appropriate to identify this type of bug since it shouldn't happen in normal execution.
Tushar7012 commentedon Feb 1, 2026
Thanks for the feedback, @Jefffrey !
I see where you're coming from, and I agree that reaching invoke_with_args() should trigger an internal error for any unexpected states. However, I was considering a small improvement for nested async UDFs, which might still be rare edge cases where an error happens that might not be immediately obvious.
The idea is to fail fast with a clear, user-facing error message as soon as we detect that we’re in an invalid nested async context, rather than letting the error bubble up as an internal error. This would just help make the failure more predictable and easier for users to diagnose — especially in cases where they’re trying to chain async UDFs in non-standard ways.
That said, I agree that we don’t want to complicate things unnecessarily, so if you feel that handling this at the invoke_with_args() level is enough (or that this extra guard isn't needed), I'm happy to skip that part and stick with the mapper changes and tests.
Let me know your thoughts!
Jefffrey commentedon Feb 1, 2026
Could you elaborate on this? How would we detect we're in an invalid state?
Tushar7012 commentedon Feb 1, 2026
Great question, @jeffrey!
By "invalid nested async context," I mean situations where we’re invoking an async UDF within another async UDF, which isn't allowed or supported. We could detect this by tracking the depth of nested async calls or using a flag to mark when we're already in an async execution.
If we hit an invalid state, we could catch it early and show a clear error instead of letting it result in a less informative internal error.
Hope this clears it up! Let me know what you think.
Jefffrey commentedon Feb 1, 2026
Aren't we fixing to support this with #20039?
Tushar7012 commentedon Feb 1, 2026
Thanks for the follow-up, @Jefffrey !
You're right — #20039 does aim to fix supporting async UDF nesting. My suggestion here is more of a precaution in case we encounter edge cases where nesting could still break unexpectedly, even after the support is added.
So, to be clear:
We’re working to support nesting with #20039, but I’m proposing an additional guard to catch and handle cases where the nesting isn’t working as expected, providing a clear error for easier debugging.
If this fits into the scope of the PR, we can include this small guard, but if we’re confident that #20039 will cover all edge cases, we can drop it and move forward with just the original fix.
Let me know your thoughts, and we can proceed!
Jefffrey commentedon Feb 1, 2026
I feel we're unnecessarily complicating this.
Tushar7012 commentedon Feb 1, 2026
Thanks for the feedback, @Jefffrey !
I agree — let’s keep it simple. Since #20039 addresses the core support for async UDF nesting, we’ll proceed with that fix alone and skip the extra guard.
Tushar7012 commentedon Feb 1, 2026
Hi @Jefffrey ,
I just wanted to check in and confirm if you're okay with moving forward with the PR as discussed. We’ll focus on #20039 for async UDF nesting support and proceed with the current fix without adding the extra guard.
Let me know if you're good to go, and I'll continue with the changes.
Thanks again!
Jefffrey commentedon Feb 2, 2026
Yes we should just proceed with fixing the problem which #20039 already aims to do
Fix Issue apache#20031: Preserve async ScalarUDF through coalescing o…
feat(rust/sedona-query-planner): skip RS_EnsureLoaded on args returni…
fix(rust/sedona-query-planner): treat columns produced by RS_EnsureLo…
fix(rust/sedona-query-planner): treat columns produced by RS_EnsureLo…