diff --git a/geoarrow-pyarrow/src/geoarrow/pyarrow/_compute.py b/geoarrow-pyarrow/src/geoarrow/pyarrow/_compute.py index d018058..ab2b506 100644 --- a/geoarrow-pyarrow/src/geoarrow/pyarrow/_compute.py +++ b/geoarrow-pyarrow/src/geoarrow/pyarrow/_compute.py @@ -251,6 +251,10 @@ def as_geoarrow(obj, type=None, coord_type=None, promote_multi=False): GeometryExtensionArray:MultiPointType(geoarrow.multipoint_z)[2] + + Large and view WKB/WKT outputs are converted through ordinary binary/string + storage and then cast to the requested storage type. Each intermediate chunk + must fit within the ordinary storage type's 32-bit offset limits. """ obj = obj_as_array_or_chunked(obj) @@ -262,6 +266,26 @@ def as_geoarrow(obj, type=None, coord_type=None, promote_multi=False): if obj.type.spec == type.spec: return obj + # geoarrow-c cannot yet write large or view WKB/WKT arrays. + if type.encoding in (Encoding.LARGE_WKB, Encoding.WKB_VIEW): + intermediate_type = _type.wkb() + elif type.encoding in (Encoding.LARGE_WKT, Encoding.WKT_VIEW): + intermediate_type = _type.wkt() + else: + intermediate_type = None + + if intermediate_type is not None: + obj = as_geoarrow(obj, intermediate_type) + storage = ensure_storage(obj).cast(type.storage_type) + # Match the kernel's behavior of preserving input extension metadata. + type = type.with_metadata(obj.type.__arrow_ext_serialize__()) + # PyArrow cannot wrap a ChunkedArray with zero chunks directly. + if isinstance(storage, pa.ChunkedArray): + return pa.chunked_array( + [type.wrap_array(chunk) for chunk in storage.chunks], type=type + ) + return type.wrap_array(storage) + return push_all(Kernel.as_geoarrow, obj, args={"type_out": type}) diff --git a/geoarrow-pyarrow/tests/test_compute.py b/geoarrow-pyarrow/tests/test_compute.py index b089b63..8e0696c 100644 --- a/geoarrow-pyarrow/tests/test_compute.py +++ b/geoarrow-pyarrow/tests/test_compute.py @@ -221,6 +221,49 @@ def test_as_geoarrow(): assert array.type.spec == ga.wkb().spec +@pytest.mark.parametrize( + "target_factory", [ga.large_wkb, ga.wkb_view, ga.large_wkt, ga.wkt_view] +) +@pytest.mark.parametrize("input_kind", ["array", "chunked", "empty", "empty_chunked"]) +def test_as_geoarrow_serialized_storage(target_factory, input_kind): + target = target_factory() + if target.encoding in (ga.Encoding.WKB_VIEW, ga.Encoding.WKT_VIEW): + if int(pa.__version__.split(".")[0]) < 18: + pytest.skip("Casting to view storage requires pyarrow >= 18") + + source = ga.as_geoarrow(["POINT (0 1)", None, "POINT (2 3)"]) + source = ga.with_crs(source, ga.OGC_CRS84) + source = ga.with_edge_type(source, ga.EdgeType.SPHERICAL) + if input_kind == "chunked": + source = pa.chunked_array([source[:1], source[1:]]) + elif input_kind == "empty": + source = source[:0] + elif input_kind == "empty_chunked": + source = pa.chunked_array([], type=source.type) + + result = ga.as_geoarrow(source, type=target) + result.validate(full=True) + assert result.type.storage_type == target.storage_type + assert ( + result.type.__arrow_ext_serialize__() == source.type.__arrow_ext_serialize__() + ) + ordinary = ( + ga.as_wkb(source) + if target.encoding in (ga.Encoding.LARGE_WKB, ga.Encoding.WKB_VIEW) + else ga.as_wkt(source) + ) + assert ( + _compute.ensure_storage(result).to_pylist() + == _compute.ensure_storage(ordinary).to_pylist() + ) + assert ga.as_geoarrow(result, type=result.type) is result + if isinstance(source, pa.ChunkedArray): + assert isinstance(result, pa.ChunkedArray) + assert [len(chunk) for chunk in result.chunks] == [ + len(chunk) for chunk in source.chunks + ] + + def test_make_point(): xs = [1, 2, 3] ys = [4, 5, 6]