Tushar7012 opened a new pull request, #20039:
URL: https://github.com/apache/datafusion/pull/20039
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases. You can
link an issue to this PR using the GitHub syntax. For example `Closes #123`
indicates that this PR will close issue #123.
-->
- Closes #20031
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in
the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your
changes and offer better suggestions for fixes.
-->
Currently, nesting asynchronous User-Defined Functions (e.g.,
`my_async_udf(my_async_udf(col))`) causes an internal error: `"async functions
should not be called directly"`. This happens because
[AsyncFuncExec](cci:2://file:///d:/Agentic_AI/Gssoc_Apache/datafusion/datafusion/physical-plan/src/async_func.rs:50:0-56:1)
and the physical planner treat all async expressions as independent and
parallelizable against the initial input batch. When one async UDF depends on
the output of another, the dependency chain is invalid, and the inner UDF's
output is not available when the outer UDF tries to execute.
This PR updates the planning and execution logic to correctly handle these
dependencies by processing async expressions sequentially when necessary and
ensuring intermediate results are available for subsequent expressions.
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is
sometimes worth providing a summary of the individual changes in this PR.
-->
- **`AsyncMapper::find_and_map`**: Updated to perform a bottom-up
transformation. It now dynamically extends a temporary schema with the output
fields of inner async functions, allowing outer functions to be correctly
planned against those intermediate columns.
- **`AsyncFuncExec::execute`**: Modified to process async expressions
incrementally. The record batch is updated with the results of each async call
before being passed to the next one, ensuring that subsequent expressions can
access the results of previous ones.
- **`AsyncFuncExec::try_new`**: Updated schema validation logic to mirror
the execution path. It now verifies expressions against an incrementally built
schema rather than validating all expressions against the static input schema.
- **`ProjectionMapping`**: Logic added to exclude columns created by inner
async UDFs from the `ProjectionMapping` validation, as these columns do not
exist in the original input schema.
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are
they covered by existing tests)?
-->
Yes. I added a new regression test
[test_nested_async_udf](cci:1://file:///d:/Agentic_AI/Gssoc_Apache/datafusion/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs:116:0-148:1)
in
[datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs](cci:7://file:///d:/Agentic_AI/Gssoc_Apache/datafusion/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs:0:0-0:0).
- The test defines a mock async UDF.
- It executes a query with nested calls: `SELECT
test_async_udf(test_async_udf(prompt)) ...`.
- It asserts that the query executes successfully and produces the correct
result (previously it panicked).
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
- **No breaking API changes.**
- **Bug Fix:** Users can now nest asynchronous UDFs in their SQL queries
without encountering internal errors.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
