james-willis opened a new pull request, #1146: URL: https://github.com/apache/sedona-db/pull/1146
## What Adds a runtime import path for natively-compiled scalar UDFs: an out-of-tree plugin can now hand SedonaDB a real `SedonaScalarKernel` (real compiled Rust, no Python callback per invocation) via a `PyCapsule`, the same way `sedona-extension`'s `SedonaCScalarKernel` ABI is already used to statically link in kernels at build time (see `c/sedona-s2geography`) -- this is the runtime counterpart of that same mechanism. Concretely: `SedonaContext.register()` gains a new protocol, `__sedonadb_native_scalar_udfs__(self) -> list[PyCapsule]`, alongside the existing `__sedonadb_internal_udf__`/`__sedonadb_internal_aggregate_udf__`/`__sedonadb_external_format__`/`__sedonadb_raster_loader__`. ## Why Today, `SedonaContext.register()`'s only path for a *scalar function* is `arrow_udf`/`sedona_scalar_udf` -- a Python callable invoked once per batch. There's no way for an out-of-tree crate (depending on `sedona-schema`/`sedona-expr` directly, like a prototype extension building a real `SedonaScalarKernel` with its own `ArgMatcher`-based dispatch) to register that kernel and get real Rust dispatch from a Python `SedonaContext` -- only a Python-callback wrapper around it. This closes that gap using the ABI SedonaDB already has, rather than inventing a new one. ## What's in it - `import_sedona_ffi_scalar_kernel` (`import_from.rs`): imports a `PyCapsule` wrapping a `SedonaCScalarKernel` into a real `ScalarKernelRef`, reading the kernel's own declared name. Mirrors `import_sedona_ffi_table_provider`'s existing pattern exactly, including the double-free-prevention `ptr::read` + `ptr::write_bytes` zeroing. - `sedona_native_scalar_udf` (`udf.rs`): a standalone pyfunction building a `PySedonaScalarUdf` from one or more same-named kernel capsules (errors on a name mismatch rather than silently registering under the wrong name) -- usable directly, the native-kernel analog of the existing `sedona_scalar_udf`. - `register_component()` (`context.rs`): the new `__sedonadb_native_scalar_udfs__` branch. A plugin's capsules can span multiple distinct function names in one call; grouped by each kernel's own declared name before registering, the same way `SedonaContext::register_scalar_kernels` already groups statically-linked kernels. - `context.py`: `__sedonadb_native_scalar_udfs__` added to `register()`'s `supported_interfaces` and docstring. ## What's NOT in it - Aggregate UDFs. `SedonaCScalarKernel` has no aggregate equivalent yet -- an accumulator's stateful lifecycle (create/update/merge/evaluate/state/size, to participate correctly in DataFusion's own parallel aggregation) needs a materially larger C ABI than a stateless scalar kernel's. This PR is scoped to the scalar case; the aggregate ABI is a separate design question. - A Python-exposed way to build an `ArgMatcher::is_extension(name)`-style matcher (`import_arg_matcher` still only recognizes the fixed literal set). Not a blocker here: a plugin's own `ArgMatcher` lives entirely inside its compiled kernel's `return_type()` and never crosses the Python boundary, so this only matters for someone trying to express that matcher from pure Python via `arrow_udf`. - A dedicated automated test for this new path specifically. `pyo3`'s `extension-module` build (required for this crate) can't embed a Python interpreter, so `Python::attach`-based `#[test]`s can't run via `cargo test` here -- confirmed directly, and consistent with this crate having zero existing Rust tests for the same reason. The underlying capsule ABI is already thoroughly tested end to end in `sedona-extension`'s own `ffi_roundtrip`/`named_kernel` tests; this PR's new code mirrors the already-proven `import_sedona_ffi_table_provider` pattern (which has the same characteristic -- no isolated unit test, only indirect coverage via a cross-context DataFrame test). Real end-to-end verification will happen naturally once an actual out-of-tree plugin exercises this path. ## Verification - `cargo check`/`clippy --all-targets -- -D warnings`/`fmt --all -- --check` for the whole workspace (excluding the environment-only `gdal-sys` bindgen issue on this machine, unrelated to this change) -- clean. - Confirmed by direct compilation that every existing `register_component` branch, `sedona_scalar_udf`, and `PySedonaScalarUdf` usage is unaffected -- this is purely additive. -- 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]
