jiayuasu commented on code in PR #937:
URL: https://github.com/apache/sedona-db/pull/937#discussion_r3391272205
##########
python/sedonadb/python/sedonadb/udf.py:
##########
@@ -314,3 +314,210 @@ def _callable_kwarg_only_names(f):
return [
k for k, p in sig.parameters.items() if p.kind ==
inspect.Parameter.KEYWORD_ONLY
]
+
+
+def arrow_aggregate_udf(
+ return_type: Any,
+ input_types: List[Union[TypeMatcher, Any]],
+ state_types: List[Any],
+ *,
+ volatility: Literal["immutable", "stable", "volatile"] = "immutable",
+ name: Optional[str] = None,
+):
+ """Decorator for Python-implemented aggregate UDFs.
+
+ Decorates a class whose instances act as a stateful accumulator. Each
+ grouped/global aggregation builds one instance per partial-state slot
+ and then merges them.
+
+ !!! warning
+ SedonaDB Python UDFs are experimental and this interface may change
+ based on user feedback.
+
+ The decorated class must define:
+
+ - `__init__(self)`: build a fresh accumulator (no arguments).
+ - `update(self, *arrays)`: receives one `pa.Array` per declared input
+ column, splatted positionally (single input today; the array can
+ contain nulls).
+ - `state(self) -> tuple`: serialize the accumulator into a tuple of
+ Python values matching `state_types` in order.
+ - `merge(self, *arrays)`: receives one `pa.Array` per element of
+ `state_types`, splatted positionally. Each array has N rows for N
+ partial states being merged in.
+ - `evaluate(self)`: return the final scalar Python value (or `None`
+ for SQL NULL) matching `return_type`.
+
+ Args:
+ return_type: A pyarrow data type for the final aggregate result.
+ Must be a concrete pyarrow type — `TypeMatcher` constants are
+ not accepted here.
+ input_types: One or more types describing the columns the
+ aggregate consumes. Each entry is either a `TypeMatcher`
+ constant (`udf.NUMERIC`, `udf.GEOMETRY`, …) or a concrete
+ pyarrow type.
+ state_types: A list of concrete pyarrow types describing the
+ serialized state. The length must match the tuple returned by
+ `state()`.
+ volatility: `"immutable"` (default), `"stable"`, or `"volatile"`.
+ name: SQL-visible name. Defaults to the decorated class name,
+ lowercased (function lookups resolve against the lowercased
+ name).
+
+ Examples:
+
+ >>> import pyarrow as pa
+ >>> import pandas as pd
+ >>> import sedonadb
+ >>> from sedonadb import udf
+ >>> sd = sedonadb.connect()
+ >>>
+ >>> @udf.arrow_aggregate_udf(
+ ... return_type=pa.float64(),
+ ... input_types=[udf.NUMERIC],
+ ... state_types=[pa.float64(), pa.int64()],
+ ... )
+ ... class my_mean:
+ ... def __init__(self):
+ ... self.total = 0.0
+ ... self.count = 0
+ ... def update(self, batch):
+ ... for v in batch:
+ ... if v.is_valid:
+ ... self.total += float(v.as_py())
+ ... self.count += 1
+ ... def state(self):
+ ... return (self.total, self.count)
+ ... def merge(self, totals, counts):
+ ... for i in range(len(totals)):
+ ... self.total += totals[i].as_py()
+ ... self.count += counts[i].as_py()
+ ... def evaluate(self):
+ ... return None if self.count == 0 else self.total / self.count
+ ...
+ >>> sd.register_udf(my_mean)
+ >>> sd.create_data_frame(
+ ... pd.DataFrame({"k": ["a", "a", "b"], "v": [1.0, 3.0, 7.0]})
+ ... ).to_view("t", overwrite=True)
+ >>> sd.sql("SELECT k, my_mean(v) AS m FROM t GROUP BY k ORDER BY
k").show()
+ ┌──────┬─────────┐
+ │ k ┆ m │
+ │ utf8 ┆ float64 │
+ ╞══════╪═════════╡
+ │ a ┆ 2.0 │
+ ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
+ │ b ┆ 7.0 │
+ └──────┴─────────┘
+ """
+
+ def decorator(cls):
+ return AggregateUdfImpl(
+ cls, return_type, input_types, state_types, volatility, name
+ )
+
+ return decorator
+
+
+class AggregateUdfImpl:
+ """Aggregate user-defined function wrapper.
+
+ Returned by [`arrow_aggregate_udf`][sedonadb.udf.arrow_aggregate_udf].
+ Holds the user's class plus the type schemas; on registration, builds
+ a factory that produces a per-accumulator `_AccumulatorWrapper`
+ bridging the user's class to the Rust `Accumulator` trait.
+ """
+
+ def __init__(
+ self,
+ user_cls,
+ return_type,
+ input_types,
+ state_types,
+ volatility: Literal["immutable", "stable", "volatile"] = "immutable",
+ name: Optional[str] = None,
+ ):
+ self._user_cls = user_cls
+ self._return_type = return_type
+ self._input_types = input_types
+ self._state_types = state_types
+ self._volatility = volatility
+ if name is None and hasattr(user_cls, "__name__"):
+ # Function lookups (con.funcs.<name> and SQL) resolve against the
+ # lowercased name, so derive a lowercased default to keep a
+ # CamelCase class reachable.
+ self._name = user_cls.__name__.lower()
Review Comment:
Done in 7f6efc1 — added a `_camel_to_snake` helper, so `MyMean` now
registers as `my_mean` (not `mymean`). Added
`test_aggregate_udf_camel_case_name` covering it.
##########
python/sedonadb/python/sedonadb/udf.py:
##########
@@ -314,3 +314,210 @@ def _callable_kwarg_only_names(f):
return [
k for k, p in sig.parameters.items() if p.kind ==
inspect.Parameter.KEYWORD_ONLY
]
+
+
+def arrow_aggregate_udf(
+ return_type: Any,
+ input_types: List[Union[TypeMatcher, Any]],
+ state_types: List[Any],
+ *,
+ volatility: Literal["immutable", "stable", "volatile"] = "immutable",
+ name: Optional[str] = None,
+):
+ """Decorator for Python-implemented aggregate UDFs.
+
+ Decorates a class whose instances act as a stateful accumulator. Each
+ grouped/global aggregation builds one instance per partial-state slot
+ and then merges them.
+
+ !!! warning
+ SedonaDB Python UDFs are experimental and this interface may change
+ based on user feedback.
+
+ The decorated class must define:
+
+ - `__init__(self)`: build a fresh accumulator (no arguments).
+ - `update(self, *arrays)`: receives one `pa.Array` per declared input
+ column, splatted positionally (single input today; the array can
+ contain nulls).
+ - `state(self) -> tuple`: serialize the accumulator into a tuple of
+ Python values matching `state_types` in order.
+ - `merge(self, *arrays)`: receives one `pa.Array` per element of
+ `state_types`, splatted positionally. Each array has N rows for N
+ partial states being merged in.
+ - `evaluate(self)`: return the final scalar Python value (or `None`
+ for SQL NULL) matching `return_type`.
+
+ Args:
+ return_type: A pyarrow data type for the final aggregate result.
+ Must be a concrete pyarrow type — `TypeMatcher` constants are
+ not accepted here.
+ input_types: One or more types describing the columns the
+ aggregate consumes. Each entry is either a `TypeMatcher`
+ constant (`udf.NUMERIC`, `udf.GEOMETRY`, …) or a concrete
+ pyarrow type.
+ state_types: A list of concrete pyarrow types describing the
+ serialized state. The length must match the tuple returned by
+ `state()`.
+ volatility: `"immutable"` (default), `"stable"`, or `"volatile"`.
+ name: SQL-visible name. Defaults to the decorated class name,
+ lowercased (function lookups resolve against the lowercased
+ name).
+
+ Examples:
+
+ >>> import pyarrow as pa
+ >>> import pandas as pd
+ >>> import sedonadb
+ >>> from sedonadb import udf
+ >>> sd = sedonadb.connect()
+ >>>
+ >>> @udf.arrow_aggregate_udf(
+ ... return_type=pa.float64(),
+ ... input_types=[udf.NUMERIC],
+ ... state_types=[pa.float64(), pa.int64()],
+ ... )
+ ... class my_mean:
+ ... def __init__(self):
+ ... self.total = 0.0
+ ... self.count = 0
+ ... def update(self, batch):
+ ... for v in batch:
+ ... if v.is_valid:
+ ... self.total += float(v.as_py())
+ ... self.count += 1
+ ... def state(self):
+ ... return (self.total, self.count)
+ ... def merge(self, totals, counts):
+ ... for i in range(len(totals)):
+ ... self.total += totals[i].as_py()
+ ... self.count += counts[i].as_py()
+ ... def evaluate(self):
+ ... return None if self.count == 0 else self.total / self.count
+ ...
+ >>> sd.register_udf(my_mean)
+ >>> sd.create_data_frame(
+ ... pd.DataFrame({"k": ["a", "a", "b"], "v": [1.0, 3.0, 7.0]})
+ ... ).to_view("t", overwrite=True)
+ >>> sd.sql("SELECT k, my_mean(v) AS m FROM t GROUP BY k ORDER BY
k").show()
+ ┌──────┬─────────┐
+ │ k ┆ m │
+ │ utf8 ┆ float64 │
+ ╞══════╪═════════╡
+ │ a ┆ 2.0 │
+ ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
+ │ b ┆ 7.0 │
+ └──────┴─────────┘
+ """
+
+ def decorator(cls):
+ return AggregateUdfImpl(
+ cls, return_type, input_types, state_types, volatility, name
+ )
+
+ return decorator
+
+
+class AggregateUdfImpl:
+ """Aggregate user-defined function wrapper.
+
+ Returned by [`arrow_aggregate_udf`][sedonadb.udf.arrow_aggregate_udf].
+ Holds the user's class plus the type schemas; on registration, builds
+ a factory that produces a per-accumulator `_AccumulatorWrapper`
+ bridging the user's class to the Rust `Accumulator` trait.
+ """
+
+ def __init__(
+ self,
+ user_cls,
+ return_type,
+ input_types,
+ state_types,
+ volatility: Literal["immutable", "stable", "volatile"] = "immutable",
+ name: Optional[str] = None,
+ ):
+ self._user_cls = user_cls
+ self._return_type = return_type
+ self._input_types = input_types
+ self._state_types = state_types
+ self._volatility = volatility
+ if name is None and hasattr(user_cls, "__name__"):
+ # Function lookups (con.funcs.<name> and SQL) resolve against the
+ # lowercased name, so derive a lowercased default to keep a
+ # CamelCase class reachable.
+ self._name = user_cls.__name__.lower()
+ else:
+ self._name = name
+
+ def __sedona_internal_udf__(self):
+ # The factory closure is invoked by the Rust kernel once per
+ # accumulator instance; it captures `self` (which only holds the
+ # user class and small type lists).
+ def factory():
+ return _AccumulatorWrapper(
+ self._user_cls(), self._return_type, self._state_types
+ )
+
+ return sedona_aggregate_udf(
+ factory,
+ self._return_type,
+ self._input_types,
+ self._state_types,
+ self._volatility,
+ self._name,
+ )
+
+
+class _AccumulatorWrapper:
+ """Bridges a user accumulator class to the Rust `Accumulator` calls.
+
+ Translates each Rust → Python crossing:
+
+ - `update(args)` / `merge(args)`: `args` is a tuple of zero-copy
+ Arrow-array handles (PySedonaValue); we convert each to a real
+ `pa.Array` and splat them positionally into the user's method.
+ - `evaluate()`: wraps the user's scalar return value as a 1-element
+ `pa.Array` of `return_type` so the Rust side can extract a
+ `ScalarValue` via the C-data interface.
+ - `state()`: wraps each scalar in the user's tuple as a 1-element
+ `pa.Array` of the matching `state_types[i]`.
+ """
+
+ def __init__(self, user_instance, return_type, state_types):
+ self._user = user_instance
+ self._return_type = return_type
+ self._state_types = state_types
+
+ @staticmethod
+ def _to_arrays(args):
+ import pyarrow as pa
+
+ # `.to_array()` forces the Array variant before the C-data import,
+ # matching the conversion the scalar-UDF examples use.
+ return [pa.array(a.to_array()) for a in args]
+
+ @staticmethod
+ def _wrap_scalar(value, pa_type):
+ import pyarrow as pa
+
+ return pa.array([value], type=pa_type)
Review Comment:
Good idea — and I confirmed `lit()` already exposes `__arrow_c_array__`, so
it'd drop into the existing import path. The wrinkle is that `lit()` infers the
type from the value, whereas `_wrap_scalar` currently coerces into the declared
`return_type`/`state_types` — so switching changes the type-handling contract
(inference vs. declared coercion). I'd like to do that as a focused follow-up
rather than fold it in here, especially since geometry returns already work
today via WKB bytes (the new `test_aggregate_udf_shapely_geometry` returns a
unioned geometry through `ga.wkb()`). Sound OK to defer?
##########
python/sedonadb/src/udf.rs:
##########
@@ -429,3 +451,235 @@ impl PySedonaValue {
)
}
}
+
+/// SedonaAggregateUdf wrapper for Python-implemented aggregate UDFs.
+///
+/// Parallel to [PySedonaScalarUdf]: holds a SedonaAggregateUDF that wraps a
+/// Python class which produces accumulator instances. Registration goes
+/// through `__sedona_internal_udf__()` on the Python side; `InternalContext`
+/// recognizes the capsule and routes to `insert_aggregate_udf`.
+#[pyclass]
+#[derive(Clone)]
+pub struct PySedonaAggregateUdf {
+ pub inner: SedonaAggregateUDF,
+}
+
+#[pymethods]
+impl PySedonaAggregateUdf {
+ fn name(&self) -> &str {
+ self.inner.name()
+ }
+
+ fn __repr__(&self) -> String {
+ self.inner.name().to_string()
+ }
+
+ fn call(&self, args: Vec<PyExpr>) -> Result<PyExpr, PySedonaError> {
+ let expr_args = args.iter().map(|arg| arg.inner.clone()).collect();
+ let aggregate_udf: AggregateUDF = self.inner.clone().into();
+ let result = aggregate_udf.call(expr_args);
+ Ok(PyExpr::new(result))
+ }
+}
+
+#[pyfunction]
+pub fn sedona_aggregate_udf<'py>(
+ py: Python<'py>,
+ py_factory: PyObject,
+ py_return_type: PyObject,
+ py_input_types: Vec<PyObject>,
+ py_state_types: Vec<PyObject>,
+ volatility: &str,
+ name: &str,
+) -> Result<PySedonaAggregateUdf, PySedonaError> {
+ let volatility = parse_volatility(volatility)?;
+
+ let arg_matchers = py_input_types
+ .iter()
+ .map(|obj| import_arg_matcher(obj.bind(py)))
+ .collect::<Result<Vec<_>, _>>()?;
+ let return_type = import_sedona_type(py_return_type.bind(py))?;
+ let matcher = ArgMatcher::new(arg_matchers, return_type);
+
+ let state_types = py_state_types
+ .iter()
+ .map(|obj| import_sedona_type(obj.bind(py)))
+ .collect::<Result<Vec<_>, _>>()?;
+
+ let kernel = PySedonaAggregateKernel {
+ matcher,
+ state_types,
+ py_factory,
+ };
+ let kernel_ref: SedonaAccumulatorRef = Arc::new(kernel);
+ let sedona_aggregate_udf = SedonaAggregateUDF::new(name, vec![kernel_ref],
volatility);
+
+ Ok(PySedonaAggregateUdf {
+ inner: sedona_aggregate_udf,
+ })
+}
+
+#[derive(Debug)]
+struct PySedonaAggregateKernel {
+ matcher: ArgMatcher,
+ state_types: Vec<SedonaType>,
+ py_factory: PyObject,
+}
+
+impl SedonaAccumulator for PySedonaAggregateKernel {
+ fn return_type(&self, args: &[SedonaType]) -> Result<Option<SedonaType>> {
+ self.matcher.match_args(args)
+ }
+
+ fn accumulator(
+ &self,
+ args: &[SedonaType],
+ output_type: &SedonaType,
+ ) -> Result<Box<dyn Accumulator>> {
+ let instance = Python::with_gil(|py| -> Result<PyObject,
PySedonaError> {
+ Ok(self.py_factory.call0(py)?)
+ })?;
+ Ok(Box::new(PySedonaAccumulator {
+ instance,
+ input_types: args.to_vec(),
+ state_types: self.state_types.clone(),
+ output_type: output_type.clone(),
+ }))
+ }
+
+ fn state_fields(&self, _args: &[SedonaType]) -> Result<Vec<FieldRef>> {
+ self.state_types
+ .iter()
+ .enumerate()
+ .map(|(i, t)|
Ok(Arc::new(t.to_storage_field(&format!("state_{i}"), true)?)))
+ .collect()
+ }
+}
+
+#[derive(Debug)]
+struct PySedonaAccumulator {
+ instance: PyObject,
+ input_types: Vec<SedonaType>,
+ state_types: Vec<SedonaType>,
+ output_type: SedonaType,
+}
+
+impl PySedonaAccumulator {
+ /// Build a tuple of PySedonaValue (Array variant) from raw Arrow arrays
+ /// for handoff into a Python method call.
+ fn arrays_to_py_values<'py>(
+ py: Python<'py>,
+ types: &[SedonaType],
+ arrays: &[ArrayRef],
+ ) -> Result<Bound<'py, PyTuple>, PySedonaError> {
+ if types.len() != arrays.len() {
+ return Err(PySedonaError::SedonaPython(format!(
+ "Internal aggregate UDF bridge: expected {} arrays, got {}",
+ types.len(),
+ arrays.len()
+ )));
+ }
+ let values: Vec<_> = zip(types, arrays)
+ .map(|(t, a)| PySedonaValue {
+ sedona_type: PySedonaType::new(t.clone()),
+ value: ColumnarValue::Array(a.clone()),
+ num_rows: a.len(),
+ })
+ .collect();
+ Ok(PyTuple::new(py, values)?)
+ }
+
+ /// Pull a ScalarValue out of a Python return value that implements
+ /// `__arrow_c_array__()` (a one-element Arrow array). The expected
+ /// SedonaType is used to validate the array's logical type.
+ fn import_scalar(
+ py: Python<'_>,
+ result: PyObject,
+ expected: &SedonaType,
+ context: &str,
+ ) -> Result<ScalarValue, PySedonaError> {
+ let result_bound = result.bind(py);
+ if !result_bound.hasattr("__arrow_c_array__")? {
+ return Err(PySedonaError::SedonaPython(format!(
+ "Expected {context} to return an object implementing
__arrow_c_array__()"
+ )));
+ }
+ let (field, array) = import_arrow_array(result_bound)?;
+ let actual = SedonaType::from_storage_field(&field)?;
+ validate_imported_type(expected, &actual, context)?;
+ if array.len() != 1 {
+ return Err(PySedonaError::SedonaPython(format!(
+ "Expected {context} to be a 1-element array; got length {}",
+ array.len()
+ )));
+ }
+ Ok(ScalarValue::try_from_array(&array, 0)?)
+ }
+}
+
+impl Accumulator for PySedonaAccumulator {
+ fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+ Python::with_gil(|py| -> Result<(), PySedonaError> {
+ let args = Self::arrays_to_py_values(py, &self.input_types,
values)?;
+ self.instance.call_method1(py, "update", (args,))?;
+ Ok(())
+ })?;
+ Ok(())
+ }
+
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ let scalar = Python::with_gil(|py| -> Result<ScalarValue,
PySedonaError> {
+ let result = self.instance.call_method0(py, "evaluate")?;
+ Self::import_scalar(py, result, &self.output_type, "aggregate UDF
evaluate()")
+ })?;
+ Ok(scalar)
+ }
+
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
Review Comment:
Agreed, leaving it for a follow-up. The commit already notes "no
GroupsAccumulator fast-path" as a known v1 limitation; I'll capture the
hold-GIL-once-per-group benefit there when we pick it up.
--
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]
