jiayuasu opened a new pull request, #937:
URL: https://github.com/apache/sedona-db/pull/937
Exposes a Python-side counterpart to `@udf.arrow_udf` for aggregates.
Bridges a Python class implementing `init` / `update` / `state` / `merge` /
`evaluate` to DataFusion's `Accumulator` plus the existing `SedonaAggregateUDF`
/ `SedonaAccumulator` Rust traits.
Next step in the UDF surface from #791. Scalar UDFs already shipped via
`@udf.arrow_udf`; this fills the aggregate gap so users can write
`df.group_by("k").agg(my_custom_udf(col("v")))`.
## API
```python
import pyarrow as pa
from sedonadb import udf
@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)
df.group_by("k").agg(m=sd.funcs.my_mean(col("v")))
sd.sql("SELECT my_mean(v) FROM t GROUP BY k")
```
### Class contract
- `update(*arrays)` — one `pa.Array` per declared input column, splatted
positionally (single input today). Array can contain nulls.
- `state() -> tuple` — Python scalars matching `state_types` in order.
- `merge(*arrays)` — one `pa.Array` per element of `state_types`, splatted
positionally. Each array has N rows for N partial states being merged.
- `evaluate() -> scalar` — Python value matching `return_type` (or `None`
for SQL NULL).
## Implementation
| File | Change |
|---|---|
| `python/sedonadb/src/udf.rs` | New `PySedonaAggregateUdf` (wraps
`SedonaAggregateUDF`) + `sedona_aggregate_udf` pyfunction constructor.
`PySedonaAggregateKernel` impls `SedonaAccumulator`, holding a Python factory
callable used to spawn a fresh accumulator per partial state.
`PySedonaAccumulator` impls DataFusion's `Accumulator`, bridging `update_batch`
/ `evaluate` / `state` / `merge_batch` through to the Python class methods.
Arrow arrays cross the boundary via the existing `PySedonaValue` (zero-copy via
`__arrow_c_array__`). |
| `python/sedonadb/src/context.rs` | `register_udf` now disambiguates scalar
vs. aggregate `__sedona_internal_udf__` return values and routes aggregate UDFs
through `insert_aggregate_udf` + `register_udaf`. |
| `python/sedonadb/python/sedonadb/udf.py` | `arrow_aggregate_udf` decorator
returns `AggregateUdfImpl` which, on registration, builds a closure-based
factory producing per-call `_AccumulatorWrapper` instances. The wrapper
converts `PySedonaValue` handoffs to `pa.Array` via the C-data interface and
wraps the user's `evaluate`/`state` return values as 1-element Arrow arrays so
the Rust side can extract `ScalarValue`s with the existing scalar-import
helper. |
| `python/sedonadb/python/sedonadb/_lib.pyi` | (Implicit via PyO3 — no
manual stubs in this package.) |
## Scope notes (follow-ups)
This is the **v1** aggregate UDF surface. Deliberately limited:
- **Single-input aggregates only.** `input_types` is wired through a single
`ArgMatcher`. Multi-input is a small extension (loop over matchers).
- **No `GroupsAccumulator` fast-path.** Grouped aggregations use
DataFusion's `Accumulator`-per-group fallback. The fast-path can land as a
separate PR once we have a use case that demands it.
- **`return_type` and `state_types` must be concrete pyarrow types.**
`TypeMatcher` constants (`udf.NUMERIC`, `udf.GEOMETRY`, …) remain valid for
`input_types` only — output types need to be concrete so the wrapper can build
1-element Arrow arrays.
## Test plan
5 tests in `python/sedonadb/tests/test_aggregate_udf.py`:
- **Global via SQL**: `SELECT my_mean(v) FROM t` — exercises just
`update_batch` + `evaluate`.
- **Grouped via DataFrame API**: `df.group_by("k").agg(...)` — exercises
`state` + `merge_batch` (slow-path Accumulator-per-group fallback).
- **Nulls in input ignored**: input array with nulls, user code checks
`v.is_valid`.
- **Empty input returns null**: zero rows → `count==0` → `evaluate` returns
`None` → SQL NULL surfaces as NaN in float64 pandas.
- **`evaluate()` wrong-type raises**: returning a string when
`return_type=float64` produces a clear error from the Rust scalar-import path.
Plus 2 doctests on the decorator itself.
Local: 5 new + 9 existing scalar UDF + 2 new doctests + `ruff format` +
`ruff check` + `cargo fmt --check` all clean.
--
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]
