jiayuasu opened a new pull request, #939:
URL: https://github.com/apache/sedona-db/pull/939
Follow-up to the aggregate UDF decorator (#937). Addresses the inline
suggestion there to use `lit()` for return values so geometry can be returned
directly.
## What changes
`evaluate()` and `state()` may now return a **Shapely / GeoPandas / pyproj
object directly** when the corresponding declared type is a geometry/geography
type — instead of requiring the user to pre-serialize to WKB bytes.
```python
@udf.arrow_aggregate_udf(
return_type=ga.wkb(), # geometry return
input_types=[udf.GEOMETRY],
state_types=[pa.binary()],
)
class geom_union:
...
def evaluate(self):
return self.acc # a Shapely geometry — no to_wkb() needed
```
## How
`_AccumulatorWrapper._wrap_scalar` now routes a **non-null value through
`lit()`** when the declared type is an Arrow **extension type**
(geometry/geography), since pyarrow can't build those from native Python
objects. Plain Arrow types — and any `None` — keep the existing
`pa.array([value], type=...)` path, so:
- int/float coercion against the declared type is unchanged,
- typed SQL NULLs still work (including null geometry via `pa.array([None],
type=ga.wkb())`),
- only the non-null geometry/geography case gains the new behavior.
`lit()` already exposes `__arrow_c_array__`, so it drops straight into the
existing Rust import path — **no Rust changes** in this PR.
## Tests
`test_aggregate_udf_shapely_geometry` now returns the unioned Shapely
geometry directly from `evaluate()` (previously WKB bytes), exercising the new
path. The remaining 12 aggregate tests are unchanged.
Local: 26 tests (13 aggregate + 13 scalar) + 2 doctests + `ruff` clean.
Python-only change.
## Not in this PR
The other deferred follow-up from #937 — a `GroupsAccumulator` fast-path
(hold the GIL once per group) — is a larger Rust change and will be its own PR.
--
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]
