paleolimbot commented on code in PR #937:
URL: https://github.com/apache/sedona-db/pull/937#discussion_r3391676419
##########
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 catch...we should have a schema request argument on `lit()` and we
don't.
--
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]
