Copilot commented on code in PR #937:
URL: https://github.com/apache/sedona-db/pull/937#discussion_r3385826604
##########
python/sedonadb/src/context.rs:
##########
@@ -17,7 +17,7 @@
use std::{collections::HashMap, sync::Arc};
use arrow_schema::DataType;
-use datafusion_expr::ScalarUDFImpl;
+use datafusion_expr::{AggregateUDFImpl, ScalarUDFImpl};
Review Comment:
`AggregateUDFImpl` and `ScalarUDFImpl` are imported but never used in this
file. With warnings denied in CI, this will fail the Rust build.
##########
python/sedonadb/src/udf.rs:
##########
@@ -21,16 +21,20 @@ use arrow_array::{
ffi::{FFI_ArrowArray, FFI_ArrowSchema},
ArrayRef,
};
-use arrow_schema::Field;
+use arrow_schema::{Field, FieldRef};
use datafusion_common::config::ConfigOptions;
use datafusion_common::{Result, ScalarValue};
-use datafusion_expr::{AggregateUDF, ColumnarValue, ScalarUDF, ScalarUDFImpl,
Volatility};
+use datafusion_expr::{
+ Accumulator, AggregateUDF, AggregateUDFImpl, ColumnarValue, ScalarUDF,
ScalarUDFImpl,
+ Volatility,
+};
Review Comment:
`AggregateUDFImpl` and `ScalarUDFImpl` are imported but never used in this
module. With warnings denied in CI, this will fail the Rust build.
##########
python/sedonadb/python/sedonadb/udf.py:
##########
@@ -314,3 +314,204 @@ 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.
+
+ 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__"):
+ self._name = user_cls.__name__
+ else:
+ self._name = name
Review Comment:
The docstring says the default SQL name is the decorated class name
lowercased, but the implementation currently keeps the class's original casing.
This is observable in `con.funcs.<name>` attribute lookup and can also affect
SQL registration/lookups depending on identifier normalization.
##########
python/sedonadb/src/udf.rs:
##########
@@ -429,3 +451,231 @@ 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,
+ ) -> Result<ScalarValue, PySedonaError> {
+ let result_bound = result.bind(py);
+ if !result_bound.hasattr("__arrow_c_array__")? {
+ return Err(PySedonaError::SedonaPython(
+ "Expected aggregate UDF Python wrapper to return an object \
+ implementing __arrow_c_array__()"
+ .to_string(),
+ ));
+ }
+ let (field, array) = import_arrow_array(result_bound)?;
+ let actual = SedonaType::from_storage_field(&field)?;
+ validate_imported_type(expected, &actual, "aggregate UDF result")?;
Review Comment:
When importing a scalar from the Python aggregate wrapper, the validation
context is always reported as "aggregate UDF result", even when the value came
from `state()` elements. A more generic context string makes the resulting
error message less misleading.
--
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]
