[PR] fix(rust/sedona-query-planner): don't wrap a raster argument that already calls an async UDF [sedona-db]

["james-willis (via GitHub)" <gi...@apache.org>]

james-willis opened a new pull request, #1380:
URL: https://github.com/apache/sedona-db/pull/1380

   ## Summary
   
   A Python raster UDF whose argument is an `RS_EnsureLoaded(...)` call fails when the UDF call is itself the raster argument of a `needs_pixels` function such as `RS_SummaryStats` or `RS_ZonalStats`:
   
   ```python
   @udf.arrow_udf(RASTER)
   def ident(rast): ...  # returns Raster.from_numpy(...) per row
   
   sd.sql("SELECT RS_SummaryStats(ident(RS_EnsureLoaded(RS_Example())), 'mean', 1)").to_arrow_table()
   # Internal error: async functions should not be called directly.
   ```
   
   The same UDF call works under a metadata function (`RS_Width(ident(RS_EnsureLoaded(...)))`), and it also works when the loading happens in a CTE first.
   
   ## Root cause
   
   `EnsureLoadedOptimizerRule` wraps every raster argument of a `needs_pixels` call in the async `RS_EnsureLoaded`. Here that argument is `ident(rs_ensureloaded(..))`, so the rewrite produces
   
   ```
   rs_summarystats(rs_ensureloaded(ident(rs_ensureloaded(RS_Example()))), 'mean', 1)
   ```
   
   which puts one async call inside another async call's arguments. DataFusion's physical planner can't hoist nested async calls (apache/datafusion#20031). Only the innermost call lands in `AsyncFuncExec`. The outer one stays in the synchronous `ProjectionExec`, and `AsyncScalarUDF::invoke_with_args` raises the internal error. `already_loaded` already skips an argument that *is* a loader, a `returns_bytes` call, or a loaded column. It doesn't cover an argument that *contains* a loader further down.
   
   ## Fix
   
   The rule no longer wraps an argument whose expression tree calls an async UDF anywhere (`calls_async_udf`: `rs_ensureloaded` by name, as `already_loaded` matches it, or any `AsyncScalarUDF`).
   
   On why this predicate is right, and not something narrower: every async call inside the wrapped subtree nests, whatever its position, whether it is `rs_ensureloaded` or another async UDF. So every such wrap fails at execution for every input, and leaving it out is never worse:
   
   - If the argument returns in-database bytes (the usual case, e.g. a Python UDF that builds its result with `Raster.from_numpy`), the kernel reads them and the query succeeds.
   - If it returns an OutDb raster, the kernel fails with its OutDb byte-access error (`RS_SummaryStats: OutDb byte access via nd_buffer() is not yet implemented; ...`), not a wrong result. The user can load the raster in a subquery or CTE instead.
   
   An argument with no async call below it is wrapped as before. If a later pass (e.g. CSE hoisting the inner loader into a projection) leaves only a column reference below, the wrap that results is not nested, and the planner hoists it normally.
   
   ## Tests
   
   - Rust unit tests in `ensure_loaded.rs`:
     - `does_not_wrap_arg_containing_a_loader` checks the expression-level rewrite.
     - `calls_async_udf_finds_any_async_udf_below` checks that the predicate matches any `AsyncScalarUDF`.
     - `optimizer_does_not_nest_a_loader_around_a_user_loader` runs the whole optimizer plus the physical planner and asserts one loader, un-nested, inside `AsyncFuncExec`.
   - Python regression test in `test_udf.py`: a Python UDF doubles the pixels of `RS_EnsureLoaded(RS_FromPath(<GeoTIFF>))` under `RS_SummaryStats`, and the test asserts the mean.
   
   The first and third Rust tests fail without the fix. The Python test fails on `main` with the original error and passes with the fix.
   


-- 
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: issues-unsubscribe@sedona.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] fix(rust/sedona-query-planner): un-nest async loaders so pixel functions over loaded rasters plan [sedona-db]

["james-willis (via GitHub)" <gi...@apache.org>]

james-willis closed pull request #1380: fix(rust/sedona-query-planner): un-nest async loaders so pixel functions over loaded rasters plan
URL: https://github.com/apache/sedona-db/pull/1380


-- 
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: issues-unsubscribe@sedona.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org