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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
