[PR] feat(rust/sedona-raster-functions): make RS_EnsureLoaded a synchronous UDF [sedona-db]

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

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

   ## Summary
   
   `RS_EnsureLoaded` becomes an ordinary synchronous scalar UDF. SedonaDB no longer registers any DataFusion async UDF, so the planner no longer has to hoist loader calls into `AsyncFuncExec`, and every shape that used to fail with `Internal error: async functions should not be called directly` now just works:
   
   ```sql
   -- a WHERE on a pixel function's output over a join (#1265)
   SELECT name, s.count FROM (
     SELECT z.name, RS_ZonalStatsAll(d.rast, z.geometry) AS s
     FROM zones z JOIN dem d ON RS_Intersects(d.rast, z.geometry)) WHERE s.count > 0;
   -- the same over a cross join, or written in an outer join's ON clause
   -- a pixel function over a metadata function over a pixel function
   SELECT RS_Value(RS_SetSRID(RS_SetValue(RS_FromPath(p), 1, 2, 2, 7), 4326), 2, 2, 1) FROM t;
   -- a Python raster UDF over an explicit loader, under a pixel function
   SELECT RS_SummaryStats(my_udf(RS_EnsureLoaded(rast)), 'mean', 1) FROM t;
   ```
   
   DataFusion plans async UDFs only in projections, filters and aggregates, cannot nest one async call inside another (apache/datafusion#20031), never hoists them out of join expressions (apache/datafusion#16520) and evaluates them for every row even inside a `CASE`. Each of those cost SedonaDB a planner workaround (#1310 merged; #1311 and #1380 open), and the next DataFusion release would have brought the next gap. A sync UDF is planned and evaluated like `RS_Width`: anywhere, nested however, lazily where the enclosing expression is lazy.
   
   This supersedes #1311 and #1380. A follow-up PR removes the workarounds that only existed for the async UDF (`sd_restore_metadata` / `WrapAsyncUdfRule`, the DataFusion-issue references in docs).
   
   ## How a sync UDF drives async loaders
   
   The loaders stay async (`AsyncRasterLoader` is unchanged, as are the GDAL, Zarr and Python bridge implementations). `invoke_with_args` spawns the load on a dedicated process-wide tokio runtime and blocks until it completes. The calling thread is usually a worker of the session's runtime; it must not drive the load itself, because once every worker waits on its own runtime nothing makes progress. The loader runtime's workers are never blocked by query execution, so the wait cannot deadlock. When the caller is a worker of a multi-threaded runtime the wait goes through `tokio::task::block_in_place`, which hands that worker's other tasks to another thread for the duration; on a current-thread runtime, or outside any runtime, the thread simply blocks. Unit tests cover all three calling contexts (the multi-thread one awaits the UDF from a task of a one-worker runtime, so it deadlocks if the wait ever blocks the worker outright).
   
   The UDF also gained a fast path: when no band of the batch is OutDb it returns its input array as is. With nesting now allowed, a redundant outer call (the planner wraps a user's explicit `RS_EnsureLoaded(..)` again under a pixel function) costs one pass over the band metadata.
   
   Trade-offs, stated plainly:
   
   - A DataFusion worker is parked for the duration of each slice's load instead of yielding. Loads were the bottleneck in those plans anyway, and `block_in_place` keeps the rest of the runtime responsive, but a current-thread runtime (ADBC) is stalled while a load is in flight.
   - Cancellation is per slice: dropping the query's stream no longer interrupts a load in progress; the GDAL loader's cooperative cancellation only fires once the slice's load completes. Slices are bounded by `sedona.raster.max_batch_bytes`, so this is bounded too.
   - The loader runtime is created on first use and, like the Zarr runtime in `sedonadb-zarr`, does not survive `fork()`.
   
   ## Byte-bounded batches without `AsyncFuncExec`
   
   `RasterBatchBudgetRule` used to swap DataFusion's `AsyncFuncExec` for `EnsureLoadedExec`, which sliced each input batch by estimated raster bytes and evaluated the async expressions per slice. There is no `AsyncFuncExec` anymore: the `rs_ensureloaded` call is evaluated by whichever `ProjectionExec`, `FilterExec` or `AggregateExec` holds it. So the rule now puts a `RasterBatchSplitExec` **under** that node. It re-slices each input batch so the estimated bytes of the rasters the node is about to load stay within the budget; the node then evaluates each slice as its own batch, which bounds the loaded pixels in flight exactly as before (one slice = one bundled loader call = one output batch; the `sedona` crate's end-to-end test for that is unchanged apart from the node name in `EXPLAIN`).
   
   Sizing follows the raster argument of every `rs_ensureloaded` call in the node:
   
   - a column argument is read as is (no rewrite of the node);
   - an expression argument (`rs_ensureloaded(RS_FromPath(path))`) is evaluated once per batch into a `__ensure_loaded_arg_k` column by a `ProjectionExec` under the split, and the call is rewritten to read that column, so `RS_FromPath` opens each file once and a volatile argument loads the raster that was sized. A `FilterExec` whose input grew this way gets a projection back to its original columns; an `AggregateExec` gets the extended input schema;
   - an expression argument in a position that is only evaluated for some rows (a `CASE` branch after the first `WHEN` condition, the right side of `AND`/`OR`) is left in place and not sized: hoisting it would evaluate it, and open its file, for rows the node never reaches. A Python test with a `CASE` over a path that does not exist pins this; it failed against a first version of this rule that hoisted everything;
   - an argument that itself contains a loader is not sized: the inner call loads it and the outer one is a pass-through.
   
   Calls that land anywhere else (a join filter, say) are evaluated unsliced. They work, where before they errored; bounding them is a possible follow-up, not a regression.
   
   The split keeps lazy evaluation: a loader inside a `CASE` branch or on the right of `AND`/`OR` is still only evaluated for the rows that reach it (the Python test with a VRT whose source file is missing pins this for a join filter). Hoisting the call out of the node, the way `AsyncFuncExec` did, would evaluate it for every row.
   
   `EXPLAIN` now shows `RasterBatchSplitExec: sized_by=[rast@0]` under the loading node instead of `EnsureLoadedExec: async_expr=[..]`. #1417 (hash-partition the loader's input) is built on the old node and needs rebasing onto this; the repartition goes under the split.
   
   ## Tests
   
   - `rs_ensure_loaded.rs`: `invoke_with_args` outside any runtime, from a current-thread runtime, and from the only worker of a multi-thread runtime; in-db input passes through untouched; a loader error comes back as an error.
   - `raster_batch_budget.rs` (rewritten for the new node, same cases plus the new ones): slicing to the budget, oversized rows and nulls, zero budget, empty batches, an expression argument evaluated once per batch (projection, filter and aggregate), expression arguments in conditional positions left in place (and a `CASE` branch no row reaches never evaluating its argument), a nested loader sized by the inner call only, a filter's pass-through batch size, untouched plans.
   - `ensure_loaded.rs`: the fixpoint test now runs against a sync stand-in and checks the physical plan carries one loader call.
   - Python, against real GeoTIFFs with exact values: `WHERE s.count > 0` over `JOIN`, `LEFT JOIN` and `CROSS JOIN`; `RS_ZonalStatsAll` in the `ON` clause of an inner and an outer join (the outer-join case was a strict xfail in #1311); a loader in an unreached `CASE` branch of a join filter is never evaluated, and the same for a `CASE` in a projection; `RS_Value(RS_SetSRID(RS_SetValue(..)))`, `RS_SummaryStats(RS_SetBandNoDataValue(RS_Clip(..)))`, a loader feeding only a non-raster argument; a Python raster UDF over `RS_EnsureLoaded` under `RS_SummaryStats` and `RS_Value`.
   


-- 
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] feat(rust/sedona-raster-functions): make RS_EnsureLoaded a synchronous UDF [sedona-db]

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

james-willis merged PR #1475:
URL: https://github.com/apache/sedona-db/pull/1475


-- 
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