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