You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@sedona.apache.org by "james-willis (via GitHub)" <gi...@apache.org> on 2026/08/10 21:22:11 UTC
[PR] Add a native scalar UDF import path for plugins [sedona-db]
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3779156049
##########
python/sedonadb/src/context.rs:
##########
@@ -354,6 +356,47 @@ impl InternalContext {
.extract::<PyRasterLoaderWrapper>()?;
self.inner.register_raster_loader(wrapper.inner);
return Ok(());
+ } else if component.hasattr("__sedonadb_native_scalar_udfs__")? {
+ // One or more natively-compiled kernel capsules (see
+ // import_sedona_ffi_scalar_kernel), potentially spanning several
+ // distinct function names -- grouped here by each kernel's own
+ // declared name, so two capsules sharing one name become one
+ // overloaded SedonaScalarUDF (proven directly: registering an
+ // Int64 kernel and a Float64 kernel under the same name and
+ // dispatching both via SQL). This grouping is scoped to this one
+ // call's own capsule list, not accumulated across calls the way
+ // SedonaContext::register_scalar_kernels accumulates
+ // statically-linked (e.g. s2geography) kernels -- registration
+ // goes through register_sedona_scalar_udf, which replaces any
+ // existing UDF of that name outright (a plain HashMap insert,
+ // same as DataFusion's own register_udf) rather than merging
+ // into it. Re-registering under a name already in use -- a
+ // plugin's own name, another plugin's, or a built-in's -- silently
+ // drops whatever was there before; this mirrors the existing,
+ // already-trusted __sedonadb_internal_udf__ path's behavior, not
+ // a new risk this protocol introduces.
+ //
+ // Volatility isn't plugin-configurable through this protocol,
+ // hardcoded to Immutable -- not because every native kernel is
+ // Immutable (RS_FromPath, for one, is Volatile), but because a
+ // bare capsule has nowhere to carry a volatility value. A plugin
+ // needing Volatile/Stable can call sedona_native_scalar_udf(...,
+ // volatility=...) directly and return the resulting
+ // PySedonaScalarUdf via the existing __sedonadb_internal_udf__
+ // protocol instead of this one.
+ let capsules = component
+ .call_method0("__sedonadb_native_scalar_udfs__")?
+ .extract::<Vec<Bound<PyAny>>>()?;
+ let mut kernels_by_name: HashMap<String, Vec<_>> = HashMap::new();
+ for capsule in &capsules {
+ let (name, kernel) = import_sedona_ffi_scalar_kernel(capsule)?;
+ kernels_by_name.entry(name).or_default().push(kernel);
+ }
+ for (name, kernels) in kernels_by_name {
+ let udf = SedonaScalarUDF::new(&name, kernels, Volatility::Immutable);
+ self.inner.register_sedona_scalar_udf(udf)?;
+ }
+ return Ok(());
Review Comment:
checked the code. append functions as a kernel-level overwrite anyway so my plan is to just adopt the "append" behavior. The only downside to this is that different implementors might have bizarrely divergent implementations but I suppose the user of the extensions should know that
--
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(python/sedonadb): Add a native scalar UDF import path for plugins [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis merged PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3769640580
##########
python/sedonadb/src/import_from.rs:
##########
@@ -211,3 +256,97 @@ pub fn check_pycapsule(obj: &Bound<PyAny>, name: &str) -> Result<*mut c_void, Py
Ok(pointer.as_ptr())
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_schema::DataType;
+ use sedona_expr::scalar_udf::SimpleSedonaScalarKernel;
+ use sedona_extension::{extension::SedonaCScalarKernel, scalar_kernel::ExportedScalarKernel};
+
+ /// A trivial real kernel (matches any single numeric arg, returns it
+ /// unchanged), exported to a `SedonaCScalarKernel` and wrapped in a real
+ /// `PyCapsule` -- the exact same export path `sedona-extension`'s own
+ /// `ffi_roundtrip`/`named_kernel` tests already prove correct end to
+ /// end. `#[dev-dependencies] pyo3 = { features = ["auto-initialize"] }`
+ /// is what makes `Python::attach` usable here at all: `extension-module`
+ /// (needed for the real wheel build) is only ever added by maturin's own
+ /// build flags, never by this crate's Cargo.toml, so it's never present
+ /// during `cargo test`.
+ fn capsule_with_named_kernel<'py>(
Review Comment:
This works too, but another test that would be good to implement is `__sedona_scalar_udf__` on our PyO3 function representation and check roundtrip behaviour.
##########
python/sedonadb/src/context.rs:
##########
@@ -354,6 +356,47 @@ impl InternalContext {
.extract::<PyRasterLoaderWrapper>()?;
self.inner.register_raster_loader(wrapper.inner);
return Ok(());
+ } else if component.hasattr("__sedonadb_native_scalar_udfs__")? {
+ // One or more natively-compiled kernel capsules (see
+ // import_sedona_ffi_scalar_kernel), potentially spanning several
+ // distinct function names -- grouped here by each kernel's own
+ // declared name, so two capsules sharing one name become one
+ // overloaded SedonaScalarUDF (proven directly: registering an
+ // Int64 kernel and a Float64 kernel under the same name and
+ // dispatching both via SQL). This grouping is scoped to this one
+ // call's own capsule list, not accumulated across calls the way
+ // SedonaContext::register_scalar_kernels accumulates
+ // statically-linked (e.g. s2geography) kernels -- registration
+ // goes through register_sedona_scalar_udf, which replaces any
+ // existing UDF of that name outright (a plain HashMap insert,
+ // same as DataFusion's own register_udf) rather than merging
+ // into it. Re-registering under a name already in use -- a
+ // plugin's own name, another plugin's, or a built-in's -- silently
+ // drops whatever was there before; this mirrors the existing,
+ // already-trusted __sedonadb_internal_udf__ path's behavior, not
+ // a new risk this protocol introduces.
+ //
+ // Volatility isn't plugin-configurable through this protocol,
+ // hardcoded to Immutable -- not because every native kernel is
+ // Immutable (RS_FromPath, for one, is Volatile), but because a
+ // bare capsule has nowhere to carry a volatility value. A plugin
+ // needing Volatile/Stable can call sedona_native_scalar_udf(...,
+ // volatility=...) directly and return the resulting
Review Comment:
We can add this later via a `get_property()` callback on the kernel
##########
python/sedonadb/python/sedonadb/context.py:
##########
@@ -490,6 +501,7 @@ def register(self, component: Any, **kwargs: Any) -> None:
"__sedonadb_internal_aggregate_udf__",
"__sedonadb_external_format__",
"__sedonadb_raster_loader__",
+ "__sedonadb_native_scalar_udfs__",
Review Comment:
The individual function definition is probably a better candidate here (i.e., `__sedonadb_scalar_udf__` or `__sedonadb_native_scalar_udf__`), mostly because we already have a function object at the Python level (e.g., `sd.funcs.st_asbinary`) that we can implement it on to ensure it works.
##########
python/sedonadb/Cargo.toml:
##########
@@ -64,3 +64,15 @@ thiserror = { workspace = true }
tokio = { workspace = true }
mimalloc = { workspace = true, optional = true }
libmimalloc-sys = { workspace = true, optional = true }
+
+[dev-dependencies]
+# `extension-module` (needed for the real cdylib/wheel build) is only ever
+# added by maturin's own build flags (see pyproject.toml's [tool.maturin]),
+# never by this Cargo.toml -- so `cargo test` builds this crate without it,
+# and `auto-initialize` here is what lets a #[test] embed and drive a real
+# Python interpreter via Python::attach. Confirmed directly: explicitly
+# compiling with `--features pyo3/extension-module` fails to link (undefined
+# libpython symbols), but plain `cargo test`/`cargo test --all-features`
+# (CI's actual invocation) never requests that feature, so the two never
+# collide in practice.
+pyo3 = { workspace = true, features = ["auto-initialize"] }
Review Comment:
As above, we should also make sure this is tested at the Python level by enabling these to get exported as well (more realistic)
##########
python/sedonadb/src/context.rs:
##########
@@ -354,6 +356,47 @@ impl InternalContext {
.extract::<PyRasterLoaderWrapper>()?;
self.inner.register_raster_loader(wrapper.inner);
return Ok(());
+ } else if component.hasattr("__sedonadb_native_scalar_udfs__")? {
+ // One or more natively-compiled kernel capsules (see
+ // import_sedona_ffi_scalar_kernel), potentially spanning several
+ // distinct function names -- grouped here by each kernel's own
+ // declared name, so two capsules sharing one name become one
+ // overloaded SedonaScalarUDF (proven directly: registering an
+ // Int64 kernel and a Float64 kernel under the same name and
+ // dispatching both via SQL). This grouping is scoped to this one
+ // call's own capsule list, not accumulated across calls the way
+ // SedonaContext::register_scalar_kernels accumulates
+ // statically-linked (e.g. s2geography) kernels -- registration
+ // goes through register_sedona_scalar_udf, which replaces any
+ // existing UDF of that name outright (a plain HashMap insert,
+ // same as DataFusion's own register_udf) rather than merging
+ // into it. Re-registering under a name already in use -- a
+ // plugin's own name, another plugin's, or a built-in's -- silently
+ // drops whatever was there before; this mirrors the existing,
+ // already-trusted __sedonadb_internal_udf__ path's behavior, not
+ // a new risk this protocol introduces.
+ //
+ // Volatility isn't plugin-configurable through this protocol,
+ // hardcoded to Immutable -- not because every native kernel is
+ // Immutable (RS_FromPath, for one, is Volatile), but because a
+ // bare capsule has nowhere to carry a volatility value. A plugin
+ // needing Volatile/Stable can call sedona_native_scalar_udf(...,
+ // volatility=...) directly and return the resulting
+ // PySedonaScalarUdf via the existing __sedonadb_internal_udf__
+ // protocol instead of this one.
+ let capsules = component
+ .call_method0("__sedonadb_native_scalar_udfs__")?
+ .extract::<Vec<Bound<PyAny>>>()?;
+ let mut kernels_by_name: HashMap<String, Vec<_>> = HashMap::new();
+ for capsule in &capsules {
+ let (name, kernel) = import_sedona_ffi_scalar_kernel(capsule)?;
+ kernels_by_name.entry(name).or_default().push(kernel);
+ }
+ for (name, kernels) in kernels_by_name {
+ let udf = SedonaScalarUDF::new(&name, kernels, Volatility::Immutable);
+ self.inner.register_sedona_scalar_udf(udf)?;
+ }
+ return Ok(());
Review Comment:
The FunctionSet for our functions already consolidates multiple kernels into one. Apologies if I missed it, but do we need to do this especially for the Python registration?
##########
python/sedonadb/src/udf.rs:
##########
@@ -207,6 +210,58 @@ pub fn sedona_scalar_udf<'py>(
})
}
+/// Build a [`PySedonaScalarUdf`] from one or more natively-compiled kernel
+/// capsules (see [`crate::import_from::import_sedona_ffi_scalar_kernel`]),
+/// instead of `sedona_scalar_udf`'s single Python-callable kernel. Real
+/// compiled Rust runs per invocation -- no GIL re-entry, no Python callback
Review Comment:
I am not sure we need this because of the `FunctionSet` registration behaviour that does this already (feel free to correct me if I missed that).
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3778623900
##########
python/sedonadb/python/sedonadb/context.py:
##########
@@ -490,6 +501,7 @@ def register(self, component: Any, **kwargs: Any) -> None:
"__sedonadb_internal_aggregate_udf__",
"__sedonadb_external_format__",
"__sedonadb_raster_loader__",
+ "__sedonadb_native_scalar_udfs__",
Review Comment:
done
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3778271752
##########
python/sedonadb/src/udf.rs:
##########
@@ -207,6 +210,58 @@ pub fn sedona_scalar_udf<'py>(
})
}
+/// Build a [`PySedonaScalarUdf`] from one or more natively-compiled kernel
+/// capsules (see [`crate::import_from::import_sedona_ffi_scalar_kernel`]),
+/// instead of `sedona_scalar_udf`'s single Python-callable kernel. Real
+/// compiled Rust runs per invocation -- no GIL re-entry, no Python callback
Review Comment:
refactored so now `sedona_native_scalar_udf` is the place that kernels are all registered at once for a `SedonaScalarUDF`
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3778622769
##########
python/sedonadb/Cargo.toml:
##########
@@ -64,3 +64,15 @@ thiserror = { workspace = true }
tokio = { workspace = true }
mimalloc = { workspace = true, optional = true }
libmimalloc-sys = { workspace = true, optional = true }
+
+[dev-dependencies]
+# `extension-module` (needed for the real cdylib/wheel build) is only ever
+# added by maturin's own build flags (see pyproject.toml's [tool.maturin]),
+# never by this Cargo.toml -- so `cargo test` builds this crate without it,
+# and `auto-initialize` here is what lets a #[test] embed and drive a real
+# Python interpreter via Python::attach. Confirmed directly: explicitly
+# compiling with `--features pyo3/extension-module` fails to link (undefined
+# libpython symbols), but plain `cargo test`/`cargo test --all-features`
+# (CI's actual invocation) never requests that feature, so the two never
+# collide in practice.
+pyo3 = { workspace = true, features = ["auto-initialize"] }
Review Comment:
done
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3778937432
##########
python/sedonadb/src/context.rs:
##########
@@ -354,6 +356,47 @@ impl InternalContext {
.extract::<PyRasterLoaderWrapper>()?;
self.inner.register_raster_loader(wrapper.inner);
return Ok(());
+ } else if component.hasattr("__sedonadb_native_scalar_udfs__")? {
+ // One or more natively-compiled kernel capsules (see
+ // import_sedona_ffi_scalar_kernel), potentially spanning several
+ // distinct function names -- grouped here by each kernel's own
+ // declared name, so two capsules sharing one name become one
+ // overloaded SedonaScalarUDF (proven directly: registering an
+ // Int64 kernel and a Float64 kernel under the same name and
+ // dispatching both via SQL). This grouping is scoped to this one
+ // call's own capsule list, not accumulated across calls the way
+ // SedonaContext::register_scalar_kernels accumulates
+ // statically-linked (e.g. s2geography) kernels -- registration
+ // goes through register_sedona_scalar_udf, which replaces any
+ // existing UDF of that name outright (a plain HashMap insert,
+ // same as DataFusion's own register_udf) rather than merging
+ // into it. Re-registering under a name already in use -- a
+ // plugin's own name, another plugin's, or a built-in's -- silently
+ // drops whatever was there before; this mirrors the existing,
+ // already-trusted __sedonadb_internal_udf__ path's behavior, not
+ // a new risk this protocol introduces.
+ //
+ // Volatility isn't plugin-configurable through this protocol,
+ // hardcoded to Immutable -- not because every native kernel is
+ // Immutable (RS_FromPath, for one, is Volatile), but because a
+ // bare capsule has nowhere to carry a volatility value. A plugin
+ // needing Volatile/Stable can call sedona_native_scalar_udf(...,
+ // volatility=...) directly and return the resulting
+ // PySedonaScalarUdf via the existing __sedonadb_internal_udf__
+ // protocol instead of this one.
+ let capsules = component
+ .call_method0("__sedonadb_native_scalar_udfs__")?
+ .extract::<Vec<Bound<PyAny>>>()?;
+ let mut kernels_by_name: HashMap<String, Vec<_>> = HashMap::new();
+ for capsule in &capsules {
+ let (name, kernel) = import_sedona_ffi_scalar_kernel(capsule)?;
+ kernels_by_name.entry(name).or_default().push(kernel);
+ }
+ for (name, kernels) in kernels_by_name {
+ let udf = SedonaScalarUDF::new(&name, kernels, Volatility::Immutable);
+ self.inner.register_sedona_scalar_udf(udf)?;
+ }
+ return Ok(());
Review Comment:
> append would silently merge a plugin's kernel into an existing same-name function.
In some cases this is what you want (e.g., add an implementation of `sd_format()` for your type to make it print nicely; add an overload of `st_intersects()` to handle your H3 type or whatever). Should there be an `overwrite=False` that defaults to the safe thing?
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3778618927
##########
python/sedonadb/src/import_from.rs:
##########
@@ -211,3 +256,97 @@ pub fn check_pycapsule(obj: &Bound<PyAny>, name: &str) -> Result<*mut c_void, Py
Ok(pointer.as_ptr())
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_schema::DataType;
+ use sedona_expr::scalar_udf::SimpleSedonaScalarKernel;
+ use sedona_extension::{extension::SedonaCScalarKernel, scalar_kernel::ExportedScalarKernel};
+
+ /// A trivial real kernel (matches any single numeric arg, returns it
+ /// unchanged), exported to a `SedonaCScalarKernel` and wrapped in a real
+ /// `PyCapsule` -- the exact same export path `sedona-extension`'s own
+ /// `ffi_roundtrip`/`named_kernel` tests already prove correct end to
+ /// end. `#[dev-dependencies] pyo3 = { features = ["auto-initialize"] }`
+ /// is what makes `Python::attach` usable here at all: `extension-module`
+ /// (needed for the real wheel build) is only ever added by maturin's own
+ /// build flags, never by this crate's Cargo.toml, so it's never present
+ /// during `cargo test`.
+ fn capsule_with_named_kernel<'py>(
Review Comment:
added the test.
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3778290690
##########
python/sedonadb/src/context.rs:
##########
@@ -354,6 +356,47 @@ impl InternalContext {
.extract::<PyRasterLoaderWrapper>()?;
self.inner.register_raster_loader(wrapper.inner);
return Ok(());
+ } else if component.hasattr("__sedonadb_native_scalar_udfs__")? {
+ // One or more natively-compiled kernel capsules (see
+ // import_sedona_ffi_scalar_kernel), potentially spanning several
+ // distinct function names -- grouped here by each kernel's own
+ // declared name, so two capsules sharing one name become one
+ // overloaded SedonaScalarUDF (proven directly: registering an
+ // Int64 kernel and a Float64 kernel under the same name and
+ // dispatching both via SQL). This grouping is scoped to this one
+ // call's own capsule list, not accumulated across calls the way
+ // SedonaContext::register_scalar_kernels accumulates
+ // statically-linked (e.g. s2geography) kernels -- registration
+ // goes through register_sedona_scalar_udf, which replaces any
+ // existing UDF of that name outright (a plain HashMap insert,
+ // same as DataFusion's own register_udf) rather than merging
+ // into it. Re-registering under a name already in use -- a
+ // plugin's own name, another plugin's, or a built-in's -- silently
+ // drops whatever was there before; this mirrors the existing,
+ // already-trusted __sedonadb_internal_udf__ path's behavior, not
+ // a new risk this protocol introduces.
+ //
+ // Volatility isn't plugin-configurable through this protocol,
+ // hardcoded to Immutable -- not because every native kernel is
+ // Immutable (RS_FromPath, for one, is Volatile), but because a
+ // bare capsule has nowhere to carry a volatility value. A plugin
+ // needing Volatile/Stable can call sedona_native_scalar_udf(...,
+ // volatility=...) directly and return the resulting
+ // PySedonaScalarUdf via the existing __sedonadb_internal_udf__
+ // protocol instead of this one.
+ let capsules = component
+ .call_method0("__sedonadb_native_scalar_udfs__")?
+ .extract::<Vec<Bound<PyAny>>>()?;
+ let mut kernels_by_name: HashMap<String, Vec<_>> = HashMap::new();
+ for capsule in &capsules {
+ let (name, kernel) = import_sedona_ffi_scalar_kernel(capsule)?;
+ kernels_by_name.entry(name).or_default().push(kernel);
+ }
+ for (name, kernels) in kernels_by_name {
+ let udf = SedonaScalarUDF::new(&name, kernels, Volatility::Immutable);
+ self.inner.register_sedona_scalar_udf(udf)?;
+ }
+ return Ok(());
Review Comment:
Same response as below about refactoring it so all kernels register together.
I kept `register_sedona_scalar_udf` (replace) rather than `add_scalar_udf_impl` (append), to match `__sedonadb_internal_udf__` and keep registration predictable — append would silently merge a plugin's kernel into an existing same-name function.
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3779785420
##########
python/sedonadb/python/sedonadb/context.py:
##########
@@ -444,6 +444,18 @@ def register(self, component: Any, **kwargs: Any) -> None:
- An ExternalFormatSpec implementing a custom datasource type
- An object implementing __sedonadb_extension__(ctx, **kwargs), which
is called with this context and any keyword arguments passed.
+ - A single function object implementing __sedonadb_scalar_udf__(self),
+ which returns a list of PyCapsule objects -- the overload kernels of
+ that one function, each wrapping a natively-compiled
+ SedonaCScalarKernel sharing the function's SQL name -- so real
+ compiled Rust runs per invocation, not a Python callback. Register
+ each function individually. The kernels are grouped into one
+ overloaded UDF; registering under a name already in use, including a
+ built-in's, replaces it, the same as __sedonadb_internal_udf__ already
+ does. Volatility is always Immutable through this protocol; a kernel
+ needing Volatile or Stable should be registered via
+ __sedonadb_internal_udf__ instead, built with
+ sedonadb._lib.sedona_native_scalar_udf(kernels, volatility=...).
Review Comment:
done
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3779783912
##########
python/sedonadb/src/context.rs:
##########
@@ -354,6 +354,23 @@ impl InternalContext {
.extract::<PyRasterLoaderWrapper>()?;
self.inner.register_raster_loader(wrapper.inner);
return Ok(());
+ } else if component.hasattr("__sedonadb_scalar_udf__")? {
+ // One function's overload kernels, each a natively-compiled kernel
+ // capsule (see import_sedona_ffi_scalar_kernel) all sharing this
+ // function's SQL name. `sedona_native_scalar_udf` imports each
+ // capsule, checks the declared names agree, and groups them into
+ // one overloaded UDF. Registration replaces any existing UDF of
+ // that name outright, mirroring the existing __sedonadb_internal_udf__
+ // path (not an append). Volatility is Immutable via this path; a
+ // plugin needing Volatile/Stable builds a PySedonaScalarUdf with
+ // sedona_native_scalar_udf(..., volatility=...) and returns it
+ // through __sedonadb_internal_udf__ instead.
Review Comment:
updated
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "james-willis (via GitHub)" <gi...@apache.org>.
james-willis commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3779786235
##########
python/sedonadb/tests/test_udf.py:
##########
@@ -218,3 +218,37 @@ def questionable_udf(arg):
match="Expected result of user-defined function to return array of length 1 but got 2",
):
con.sql("SELECT questionable_udf(123) as col").to_pandas()
+
+
+def test_native_scalar_udf_export_import_roundtrip(con):
+ # A SedonaDB built-in scalar function exposes its native overload kernels
+ # via __sedonadb_scalar_udf__ (the export side of the plugin protocol).
+ # Exporting those capsules and rebuilding a UDF under a new name must
+ # produce a function whose output is identical to the built-in's, proving
+ # the native kernels survive a full export -> capsule -> import roundtrip.
+ from sedonadb._lib import sedona_native_scalar_udf
+
+ st_asbinary = con.funcs.st_asbinary._impl
+ assert hasattr(st_asbinary, "__sedonadb_scalar_udf__")
Review Comment:
done
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3779010960
##########
python/sedonadb/src/context.rs:
##########
@@ -354,6 +354,23 @@ impl InternalContext {
.extract::<PyRasterLoaderWrapper>()?;
self.inner.register_raster_loader(wrapper.inner);
return Ok(());
+ } else if component.hasattr("__sedonadb_scalar_udf__")? {
+ // One function's overload kernels, each a natively-compiled kernel
+ // capsule (see import_sedona_ffi_scalar_kernel) all sharing this
+ // function's SQL name. `sedona_native_scalar_udf` imports each
+ // capsule, checks the declared names agree, and groups them into
+ // one overloaded UDF. Registration replaces any existing UDF of
+ // that name outright, mirroring the existing __sedonadb_internal_udf__
+ // path (not an append). Volatility is Immutable via this path; a
+ // plugin needing Volatile/Stable builds a PySedonaScalarUdf with
+ // sedona_native_scalar_udf(..., volatility=...) and returns it
+ // through __sedonadb_internal_udf__ instead.
Review Comment:
An earlier comment suggested we don't clobber existing names. Does this comment need to be updated?
##########
python/sedonadb/python/sedonadb/context.py:
##########
@@ -444,6 +444,18 @@ def register(self, component: Any, **kwargs: Any) -> None:
- An ExternalFormatSpec implementing a custom datasource type
- An object implementing __sedonadb_extension__(ctx, **kwargs), which
is called with this context and any keyword arguments passed.
+ - A single function object implementing __sedonadb_scalar_udf__(self),
+ which returns a list of PyCapsule objects -- the overload kernels of
+ that one function, each wrapping a natively-compiled
+ SedonaCScalarKernel sharing the function's SQL name -- so real
+ compiled Rust runs per invocation, not a Python callback. Register
+ each function individually. The kernels are grouped into one
+ overloaded UDF; registering under a name already in use, including a
+ built-in's, replaces it, the same as __sedonadb_internal_udf__ already
+ does. Volatility is always Immutable through this protocol; a kernel
+ needing Volatile or Stable should be registered via
+ __sedonadb_internal_udf__ instead, built with
+ sedonadb._lib.sedona_native_scalar_udf(kernels, volatility=...).
Review Comment:
It's a little strange to reference a method from _lib here. Can you add this as a regular function in `udf.py`?
##########
python/sedonadb/tests/test_udf.py:
##########
@@ -218,3 +218,37 @@ def questionable_udf(arg):
match="Expected result of user-defined function to return array of length 1 but got 2",
):
con.sql("SELECT questionable_udf(123) as col").to_pandas()
+
+
+def test_native_scalar_udf_export_import_roundtrip(con):
+ # A SedonaDB built-in scalar function exposes its native overload kernels
+ # via __sedonadb_scalar_udf__ (the export side of the plugin protocol).
+ # Exporting those capsules and rebuilding a UDF under a new name must
+ # produce a function whose output is identical to the built-in's, proving
+ # the native kernels survive a full export -> capsule -> import roundtrip.
+ from sedonadb._lib import sedona_native_scalar_udf
+
+ st_asbinary = con.funcs.st_asbinary._impl
+ assert hasattr(st_asbinary, "__sedonadb_scalar_udf__")
Review Comment:
Can you ensure this is also present on whatever object is returned by `con.funcs.st_asbinary` to avoid the `_impl` hop?
--
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] Add a native scalar UDF import path for plugins [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #1146:
URL: https://github.com/apache/sedona-db/pull/1146#discussion_r3779281079
##########
python/sedonadb/src/context.rs:
##########
@@ -354,6 +356,47 @@ impl InternalContext {
.extract::<PyRasterLoaderWrapper>()?;
self.inner.register_raster_loader(wrapper.inner);
return Ok(());
+ } else if component.hasattr("__sedonadb_native_scalar_udfs__")? {
+ // One or more natively-compiled kernel capsules (see
+ // import_sedona_ffi_scalar_kernel), potentially spanning several
+ // distinct function names -- grouped here by each kernel's own
+ // declared name, so two capsules sharing one name become one
+ // overloaded SedonaScalarUDF (proven directly: registering an
+ // Int64 kernel and a Float64 kernel under the same name and
+ // dispatching both via SQL). This grouping is scoped to this one
+ // call's own capsule list, not accumulated across calls the way
+ // SedonaContext::register_scalar_kernels accumulates
+ // statically-linked (e.g. s2geography) kernels -- registration
+ // goes through register_sedona_scalar_udf, which replaces any
+ // existing UDF of that name outright (a plain HashMap insert,
+ // same as DataFusion's own register_udf) rather than merging
+ // into it. Re-registering under a name already in use -- a
+ // plugin's own name, another plugin's, or a built-in's -- silently
+ // drops whatever was there before; this mirrors the existing,
+ // already-trusted __sedonadb_internal_udf__ path's behavior, not
+ // a new risk this protocol introduces.
+ //
+ // Volatility isn't plugin-configurable through this protocol,
+ // hardcoded to Immutable -- not because every native kernel is
+ // Immutable (RS_FromPath, for one, is Volatile), but because a
+ // bare capsule has nowhere to carry a volatility value. A plugin
+ // needing Volatile/Stable can call sedona_native_scalar_udf(...,
+ // volatility=...) directly and return the resulting
+ // PySedonaScalarUdf via the existing __sedonadb_internal_udf__
+ // protocol instead of this one.
+ let capsules = component
+ .call_method0("__sedonadb_native_scalar_udfs__")?
+ .extract::<Vec<Bound<PyAny>>>()?;
+ let mut kernels_by_name: HashMap<String, Vec<_>> = HashMap::new();
+ for capsule in &capsules {
+ let (name, kernel) = import_sedona_ffi_scalar_kernel(capsule)?;
+ kernels_by_name.entry(name).or_default().push(kernel);
+ }
+ for (name, kernels) in kernels_by_name {
+ let udf = SedonaScalarUDF::new(&name, kernels, Volatility::Immutable);
+ self.inner.register_sedona_scalar_udf(udf)?;
+ }
+ return Ok(());
Review Comment:
It's a good point that we can/should provide some safety around that at some point...the existing system was designed to optimize for flexibility (and explicitly support overlapping definitions to enable multiple implementations of st_intersects by various libraries while we benchmarked them). Extensions can check this themselves too (via `"foofy" in ctx.funcs` or similar).
--
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