You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@sedona.apache.org by "jiayuasu (via GitHub)" <gi...@apache.org> on 2026/06/11 17:46:58 UTC
[PR] feat(python/sedonadb): let aggregate UDFs return geometry via lit() [sedona-db]
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): let aggregate UDFs return geometry via lit() [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu commented on PR #939:
URL: https://github.com/apache/sedona-db/pull/939#issuecomment-4687251077
Makes sense — agreed, and no need to apologize. The WKB-bytes geometry return already works on `main` (shipped in #937) with full metadata since `pa.array([wkb], type=...)` honors the declared CRS/edges; this PR was only adding the Shapely-object shortcut, which can't carry the declared type's metadata through `lit()` today. Closing this and we can revisit cleanly once `lit()` accepts a type request.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): let aggregate UDFs return geometry via lit() [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu commented on code in PR #939:
URL: https://github.com/apache/sedona-db/pull/939#discussion_r3399277116
##########
python/sedonadb/tests/test_aggregate_udf.py:
##########
@@ -325,13 +325,16 @@ def update(self, batch):
self._absorb(batch.to_pylist())
def state(self):
+ # State serializes to WKB binary (a plain Arrow type).
return (None if self.acc is None else shapely.to_wkb(self.acc),)
def merge(self, states):
self._absorb(states.to_pylist())
def evaluate(self):
- return None if self.acc is None else shapely.to_wkb(self.acc)
+ # Return the Shapely geometry directly; the wrapper routes it
+ # through lit() since the declared return type is an extension.
+ return self.acc
Review Comment:
Added in 10fd6e6f `test_aggregate_udf_geometry_wkb_bytes_preserves_crs` — an aggregate that returns already-serialized WKB bytes (not a Shapely object), with a declared `ga.wkb().with_crs('EPSG:4326')` return type. It asserts the output SRID is 4326, so it covers both the WKB-bytes return path and the metadata-preservation point below.
##########
python/sedonadb/python/sedonadb/udf.py:
##########
@@ -512,6 +515,16 @@ def _to_arrays(args):
def _wrap_scalar(value, pa_type):
import pyarrow as pa
+ # Extension types (e.g. geometry/geography) can't be built by pyarrow
+ # from native Python objects like Shapely geometries, so route a
+ # non-null value through `lit()`, which understands Shapely / GeoPandas
+ # / pyproj inputs. Plain Arrow types — and any null — go through
+ # `pa.array` with the declared type, preserving int/float coercion and
+ # producing a correctly typed null.
+ if value is not None and isinstance(pa_type, pa.ExtensionType):
+ from sedonadb.expr import lit
+
+ return lit(value)
return pa.array([value], type=pa_type)
Review Comment:
Good catch — and it was a regression vs. the pre-PR baseline (the original `pa.array([value], type=pa_type)` preserved CRS/edges for WKB-bytes returns). Fixed in 10fd6e6f exactly as you suggested: within the extension branch, build with the declared type first (preserving CRS/edge metadata), and only fall back to `lit()` when pyarrow can't build it (the Shapely-object case). I scoped the try/except to extension types so a plain-type mismatch (e.g. a str for float64) still surfaces pyarrow's clear "Could not convert" error rather than falling through to lit(). Verified end-to-end: a WKB-bytes return with a declared-CRS type now round-trips SRID 4326 (the new test).
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): let aggregate UDFs return geometry via lit() [sedona-db]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #939:
URL: https://github.com/apache/sedona-db/pull/939#discussion_r3398246822
##########
python/sedonadb/tests/test_aggregate_udf.py:
##########
@@ -325,13 +325,16 @@ def update(self, batch):
self._absorb(batch.to_pylist())
def state(self):
+ # State serializes to WKB binary (a plain Arrow type).
return (None if self.acc is None else shapely.to_wkb(self.acc),)
def merge(self, states):
self._absorb(states.to_pylist())
def evaluate(self):
- return None if self.acc is None else shapely.to_wkb(self.acc)
+ # Return the Shapely geometry directly; the wrapper routes it
+ # through lit() since the declared return type is an extension.
+ return self.acc
Review Comment:
This test now exercises the new Shapely-return path, but it no longer covers the previously-supported behavior where an aggregate UDF with `return_type=ga.wkb()` returns already-serialized WKB bytes from `evaluate()`. Given the wrapper still aims to accept both forms, it’s worth keeping a regression assertion for the WKB-bytes return case as well.
##########
python/sedonadb/python/sedonadb/udf.py:
##########
@@ -512,6 +515,16 @@ def _to_arrays(args):
def _wrap_scalar(value, pa_type):
import pyarrow as pa
+ # Extension types (e.g. geometry/geography) can't be built by pyarrow
+ # from native Python objects like Shapely geometries, so route a
+ # non-null value through `lit()`, which understands Shapely / GeoPandas
+ # / pyproj inputs. Plain Arrow types — and any null — go through
+ # `pa.array` with the declared type, preserving int/float coercion and
+ # producing a correctly typed null.
+ if value is not None and isinstance(pa_type, pa.ExtensionType):
+ from sedonadb.expr import lit
+
+ return lit(value)
return pa.array([value], type=pa_type)
Review Comment:
For extension-typed outputs, returning `lit(value)` ignores the declared `pa_type`. That means any extension metadata encoded in the declared type (e.g., CRS / edge_type on geoarrow geometry types) is not preserved, and can cause type-validation failures when `return_type`/`state_types` specify non-default metadata.
A more robust approach is to first try building a 1-element array with the declared type (`pa.array([value], type=pa_type)`), and only fall back to a `lit()`-style conversion when that fails (e.g., for Shapely objects). In the fallback, consider wrapping the resolved extension array’s storage with the declared `pa_type` (when supported) so the declared metadata is preserved.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): let aggregate UDFs return geometry via lit() [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu commented on PR #939:
URL: https://github.com/apache/sedona-db/pull/939#issuecomment-4687251251
Deferring the Shapely-object aggregate-UDF return path until `lit()` supports a type request (so declared geometry metadata like CRS is preserved). WKB-bytes geometry returns already work on main.
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): let aggregate UDFs return geometry via lit() [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu closed pull request #939: feat(python/sedonadb): let aggregate UDFs return geometry via lit()
URL: https://github.com/apache/sedona-db/pull/939
--
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: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org