jiayuasu commented on code in PR #937:
URL: https://github.com/apache/sedona-db/pull/937#discussion_r3385866624
##########
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
+
+ def __sedona_internal_udf__(self):
+ # Capture the user class and the schemas in a factory closure that
+ # the Rust kernel invokes once per accumulator instance.
+ return_type = self._return_type
+ state_types = self._state_types
+ user_cls = self._user_cls
+
+ def factory():
+ return _AccumulatorWrapper(user_cls(), return_type, 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
+
+ def update(self, args):
+ import pyarrow as pa
+
+ arrays = [pa.array(a) for a in args]
+ self._user.update(*arrays)
Review Comment:
Fixed in fa2a55f. Both `update` and `merge` go through a shared `_to_arrays`
helper which now does `pa.array(a.to_array())` — forcing the Array variant
before the C-data import, matching the scalar-UDF examples.
##########
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
+
+ def __sedona_internal_udf__(self):
+ # Capture the user class and the schemas in a factory closure that
+ # the Rust kernel invokes once per accumulator instance.
+ return_type = self._return_type
+ state_types = self._state_types
+ user_cls = self._user_cls
+
+ def factory():
+ return _AccumulatorWrapper(user_cls(), return_type, 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
+
+ def update(self, args):
+ import pyarrow as pa
+
+ arrays = [pa.array(a) for a in args]
+ self._user.update(*arrays)
+
+ def merge(self, args):
+ import pyarrow as pa
+
+ arrays = [pa.array(a) for a in args]
+ self._user.merge(*arrays)
Review Comment:
Fixed in fa2a55f. Both `update` and `merge` go through a shared `_to_arrays`
helper which now does `pa.array(a.to_array())` — forcing the Array variant
before the C-data import, matching the scalar-UDF examples.
##########
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:
Fixed in fa2a55f — the auto-derived default name is now lowercased
(`user_cls.__name__.lower()`), so the implementation matches the docstring and
a CamelCase class stays reachable via the lowercasing lookup path. Docstring
also clarified.
##########
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:
Fixed in fa2a55f — the auto-derived default name is now lowercased
(`user_cls.__name__.lower()`), so the implementation matches the docstring and
a CamelCase class stays reachable via the lowercasing lookup path. Docstring
also clarified.
##########
python/sedonadb/tests/test_aggregate_udf.py:
##########
@@ -0,0 +1,125 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pandas as pd
+import pandas.testing as pdt
+import pyarrow as pa
+import pytest
+from sedonadb import udf
+from sedonadb.expr import col
+
+
+def _mean_class():
+ @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
+
+ return my_mean
+
+
+def test_aggregate_udf_global_via_sql(con):
+ con.register_udf(_mean_class())
+ con.create_data_frame(pd.DataFrame({"v": [1.0, 3.0, 5.0, 7.0]})).to_view(
+ "t_global", overwrite=True
+ )
+ out = con.sql("SELECT my_mean(v) AS m FROM t_global").to_pandas()
+ pdt.assert_frame_equal(out, pd.DataFrame({"m": [4.0]}))
+
+
+def test_aggregate_udf_grouped(con):
+ # Grouped aggregation exercises both update_batch and the merge path
+ # across groups under the slow Accumulator-per-group fallback.
+ con.register_udf(_mean_class())
+ df = con.create_data_frame(
+ pd.DataFrame({"k": ["a", "a", "b", "b", "b"], "v": [1.0, 3.0, 2.0,
4.0, 6.0]})
+ )
+ out =
df.group_by("k").agg(m=con.funcs.my_mean(col("v"))).sort("k").to_pandas()
+ pdt.assert_frame_equal(out, pd.DataFrame({"k": ["a", "b"], "m": [2.0,
4.0]}))
+
+
+def test_aggregate_udf_input_nulls_ignored(con):
+ con.register_udf(_mean_class())
+ df = con.create_data_frame(pd.DataFrame({"v": [1.0, None, 3.0, None,
5.0]}))
+ out = df.agg(m=con.funcs.my_mean(col("v"))).to_pandas()
+ pdt.assert_frame_equal(out, pd.DataFrame({"m": [3.0]}))
+
+
+def test_aggregate_udf_empty_input_returns_null(con):
+ # No rows → count==0 → evaluate returns None → SQL NULL surfaces as
+ # NaN in float64 pandas.
+ con.register_udf(_mean_class())
+ con.create_data_frame(pd.DataFrame({"v": pd.Series([],
dtype="float64")})).to_view(
+ "t_empty", overwrite=True
+ )
+ out = con.sql("SELECT my_mean(v) AS m FROM t_empty").to_pandas()
+ assert len(out) == 1
+ assert pd.isna(out["m"].iloc[0])
+
+
+def test_aggregate_udf_evaluate_returning_wrong_type_raises(con):
+ # evaluate() must return a value compatible with return_type;
+ # returning a string when float64 is declared surfaces a clear error
+ # from the Rust scalar-import path.
+ @udf.arrow_aggregate_udf(
+ return_type=pa.float64(),
+ input_types=[udf.NUMERIC],
+ state_types=[pa.int64()],
+ )
+ class wrong_return_type:
+ def __init__(self):
+ self.n = 0
+
+ def update(self, batch):
+ self.n += len(batch)
+
+ def state(self):
+ return (self.n,)
+
+ def merge(self, counts):
+ for i in range(len(counts)):
+ self.n += counts[i].as_py()
+
+ def evaluate(self):
+ return "not a number"
+
+ con.register_udf(wrong_return_type)
+ df = con.create_data_frame(pd.DataFrame({"v": [1.0, 2.0, 3.0]}))
+ with pytest.raises(Exception):
+ df.agg(n=con.funcs.wrong_return_type(col("v"))).to_pandas()
Review Comment:
Fixed in fa2a55f — now `pytest.raises(Exception, match="Could not
convert")`. The wrong-type value fails in pyarrow when the wrapper builds the
1-element result array (`ArrowInvalid: Could not convert 'not a number' ... to
double`); updated the comment to reflect that path.
##########
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:
Good catch — fixed in fa2a55f. `import_scalar` now takes a `context`
parameter; `evaluate()` passes "aggregate UDF evaluate()" and `state()` passes
"aggregate UDF state() element", so the error message names the right source.
--
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]
