Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions geoarrow-pandas/src/geoarrow/pandas/lib.py
Original file line number Diff line number Diff line change
Expand Up @@ -507,18 +507,18 @@ def bounds(self):
"""See :func:`geoarrow.pyarrow.box`"""
array_or_chunked = _ga.box(self._obj)
if isinstance(array_or_chunked, _pa.ChunkedArray):
flattened = [chunk.flatten() for chunk in array_or_chunked.chunks]
flattened = [chunk.storage.flatten() for chunk in array_or_chunked.chunks]
seriesish = [
_pa.chunked_array(item, _pa.float64()) for item in zip(*flattened)
]
else:
seriesish = array_or_chunked.flatten()
seriesish = array_or_chunked.storage.flatten()

return _pd.DataFrame(
{
"xmin": seriesish[0],
"xmax": seriesish[1],
"ymin": seriesish[2],
"xmax": seriesish[2],
"ymin": seriesish[1],
"ymax": seriesish[3],
},
index=self._obj.index,
Expand Down
49 changes: 9 additions & 40 deletions geoarrow-pyarrow/src/geoarrow/pyarrow/_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,49 +78,18 @@ def __repr__(self):
return f"{type_name}:{repr(self.type)}[{len(self)}]\n{items_str}".strip()


class PointArray(GeometryExtensionArray):
pass


class LinestringArray(GeometryExtensionArray):
pass


class PolygonArray(GeometryExtensionArray):
pass


class MultiPointArray(GeometryExtensionArray):
pass


class MultiLinestringArray(GeometryExtensionArray):
pass


class MultiPolygonArray(GeometryExtensionArray):
pass
class BoxArray(GeometryExtensionArray):
def __repr__(self):
type_name = type(self).__name__
items_str = "\n".join(repr(item.bounds) for item in self)
return f"{type_name}:{repr(self.type)}[{len(self)}]\n{items_str}".strip()


def array_cls_from_name(name):
if name == "geoarrow.wkb":
return GeometryExtensionArray
elif name == "geoarrow.wkt":
return GeometryExtensionArray
elif name == "geoarrow.point":
return PointArray
elif name == "geoarrow.linestring":
return LinestringArray
elif name == "geoarrow.polygon":
return PolygonArray
elif name == "geoarrow.multipoint":
return MultiPointArray
elif name == "geoarrow.multilinestring":
return MultiLinestringArray
elif name == "geoarrow.multipolygon":
return MultiPolygonArray
if name == "geoarrow.box":
return BoxArray
else:
raise ValueError(f'Expected valid extension name but got "{name}"')
return GeometryExtensionArray


# Inject array_cls_from_name exactly once to avoid circular import
Expand All @@ -142,7 +111,7 @@ def array(obj, type_=None, *args, **kwargs) -> GeometryExtensionArray:
GeometryExtensionArray:WktType(geoarrow.wkt)[1]
<POINT (0 1)>
>>> ga.as_geoarrow(["POINT (0 1)"])
PointArray:PointType(geoarrow.point)[1]
GeometryExtensionArray:PointType(geoarrow.point)[1]
<POINT (0 1)>
"""
# Convert GeoPandas to WKB
Expand Down
39 changes: 20 additions & 19 deletions geoarrow-pyarrow/src/geoarrow/pyarrow/_compute.py
Original file line number Diff line number Diff line change
Expand Up @@ -251,7 +251,7 @@ def as_geoarrow(obj, type=None, coord_type=None, promote_multi=False):

>>> import geoarrow.pyarrow as ga
>>> ga.as_geoarrow(["POINT (0 1)", "MULTIPOINT Z (0 1 2, 4 5 6)"])
MultiPointArray:MultiPointType(geoarrow.multipoint_z)[2]
GeometryExtensionArray:MultiPointType(geoarrow.multipoint_z)[2]
<MULTIPOINT Z (0 1 nan)>
<MULTIPOINT Z (0 1 2, 4 5 6)>
"""
Expand Down Expand Up @@ -307,7 +307,7 @@ def make_point(x, y, z=None, m=None, crs=None):

>>> import geoarrow.pyarrow as ga
>>> ga.make_point([1, 2, 3], [4, 5, 6])
PointArray:PointType(geoarrow.point)[3]
GeometryExtensionArray:PointType(geoarrow.point)[3]
<POINT (1 4)>
<POINT (2 5)>
<POINT (3 6)>
Expand Down Expand Up @@ -338,10 +338,11 @@ def make_point(x, y, z=None, m=None, crs=None):

def _box_point_struct(storage):
arrays = storage.flatten()
return pa.StructArray.from_arrays(
[arrays[0], arrays[0], arrays[1], arrays[1]],
names=["xmin", "xmax", "ymin", "ymax"],
box_storage = pa.StructArray.from_arrays(
[arrays[0], arrays[1], arrays[0], arrays[1]],
names=["xmin", "ymin", "xmax", "ymax"],
)
return _type.types.box().to_pyarrow().wrap_array(box_storage)


def box(obj):
Expand All @@ -350,7 +351,7 @@ def box(obj):

>>> import geoarrow.pyarrow as ga
>>> ga.box(["LINESTRING (0 10, 34 -1)"]).type
StructType(struct<xmin: double, xmax: double, ymin: double, ymax: double>)
BoxType(geoarrow.box)
>>> print(str(ga.box(["LINESTRING (0 10, 34 -1)"])))
-- is_valid: all not null
-- child 0 type: double
Expand All @@ -359,11 +360,11 @@ def box(obj):
]
-- child 1 type: double
[
34
-1
]
-- child 2 type: double
[
-1
34
]
-- child 3 type: double
[
Expand Down Expand Up @@ -399,15 +400,15 @@ def _box_agg_point_struct(arrays):
out = [list(pc.min_max(array).values()) for array in arrays]
out_dict = {
"xmin": out[0][0].as_py(),
"xmax": out[0][1].as_py(),
"ymin": out[1][0].as_py(),
"xmax": out[0][1].as_py(),
"ymax": out[1][1].as_py(),
}

# Apparently pyarrow reorders dict keys when inferring scalar types?
return pa.scalar(
out_dict, pa.struct([(nm, pa.float64()) for nm in out_dict.keys()])
)
storage_type = pa.struct([(nm, pa.float64()) for nm in out_dict.keys()])
storage_array = pa.array([out_dict], storage_type)
return _type.types.box().to_pyarrow().wrap_array(storage_array)[0]


def box_agg(obj):
Expand All @@ -417,7 +418,7 @@ def box_agg(obj):

>>> import geoarrow.pyarrow as ga
>>> ga.box_agg(["POINT (0 10)", "POINT (34 -1)"])
<pyarrow.StructScalar: [('xmin', 0.0), ('xmax', 34.0), ('ymin', -1.0), ('ymax', 10.0)]>
BoxScalar({'xmin': 0.0, 'ymin': -1.0, 'xmax': 34.0, 'ymax': 10.0})
"""

obj = obj_as_array_or_chunked(obj)
Expand Down Expand Up @@ -495,7 +496,7 @@ def with_coord_type(obj, coord_type):

>>> import geoarrow.pyarrow as ga
>>> ga.with_coord_type(["POINT (0 1)"], ga.CoordType.INTERLEAVED)
PointArray:PointType(interleaved geoarrow.point)[1]
GeometryExtensionArray:PointType(interleaved geoarrow.point)[1]
<POINT (0 1)>
"""
return as_geoarrow(obj, coord_type=coord_type)
Expand Down Expand Up @@ -537,10 +538,10 @@ def with_dimensions(obj, dimensions):

>>> import geoarrow.pyarrow as ga
>>> ga.with_dimensions(["POINT (0 1)"], ga.Dimensions.XYZM)
PointArray:PointType(geoarrow.point_zm)[1]
GeometryExtensionArray:PointType(geoarrow.point_zm)[1]
<POINT ZM (0 1 nan nan)>
>>> ga.with_dimensions(["POINT ZM (0 1 2 3)"], ga.Dimensions.XY)
PointArray:PointType(geoarrow.point)[1]
GeometryExtensionArray:PointType(geoarrow.point)[1]
<POINT (0 1)>
"""
obj = as_geoarrow(obj)
Expand All @@ -557,13 +558,13 @@ def with_geometry_type(obj, geometry_type):

>>> import geoarrow.pyarrow as ga
>>> ga.with_geometry_type(["POINT (0 1)"], ga.GeometryType.MULTIPOINT)
MultiPointArray:MultiPointType(geoarrow.multipoint)[1]
GeometryExtensionArray:MultiPointType(geoarrow.multipoint)[1]
<MULTIPOINT (0 1)>
>>> ga.with_geometry_type(["MULTIPOINT (0 1)"], ga.GeometryType.POINT)
PointArray:PointType(geoarrow.point)[1]
GeometryExtensionArray:PointType(geoarrow.point)[1]
<POINT (0 1)>
>>> ga.with_geometry_type(["LINESTRING EMPTY", "POINT (0 1)"], ga.GeometryType.POINT)
PointArray:PointType(geoarrow.point)[2]
GeometryExtensionArray:PointType(geoarrow.point)[2]
<POINT (nan nan)>
<POINT (0 1)>
>>> ga.with_geometry_type(["MULTIPOINT (0 1, 2 3)"], ga.GeometryType.POINT)
Expand Down
52 changes: 47 additions & 5 deletions geoarrow-pyarrow/src/geoarrow/pyarrow/_kernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,21 +2,29 @@

import pyarrow as pa
import pyarrow_hotfix as _ # noqa: F401
from geoarrow.types import box as box_spec
from geoarrow.pyarrow._type import GeometryExtensionType


_lazy_lib = None
_geoarrow_c_version = None


def _geoarrow_c():
global _lazy_lib
global _lazy_lib, _geoarrow_c_version
if _lazy_lib is None:
try:
from geoarrow.c import lib
import geoarrow.c

except ImportError as e:
raise ImportError("Requested operation requires geoarrow-c") from e

_lazy_lib = lib
_lazy_lib = geoarrow.c.lib
if hasattr(geoarrow.c, "__version_tuple__"):
_geoarrow_c_version = geoarrow.c.__version_tuple__
else:
_geoarrow_c_version = (0, 1, 0)

return _lazy_lib


Expand Down Expand Up @@ -109,11 +117,19 @@ def unique_geometry_types_agg(type_in):

@staticmethod
def box(type_in):
return Kernel("box", type_in)
kernel = Kernel("box", type_in)
if _geoarrow_c_version <= (0, 1, 3):
return BoxKernelCompat(kernel)
else:
return kernel

@staticmethod
def box_agg(type_in):
return Kernel("box_agg", type_in)
kernel = Kernel("box_agg", type_in)
if _geoarrow_c_version <= (0, 1, 3):
return BoxKernelCompat(kernel)
else:
return kernel

@staticmethod
def _pack_options(options):
Expand All @@ -132,3 +148,29 @@ def _pack_options(options):
bytes += v.encode("UTF-8")

return bytes


class BoxKernelCompat:
"""A wrapper around the "box" kernel that works for geoarrow-c 0.1.
This is mostly to ease the transition for geoarrow-python CI while
all the packages are being updated."""

def __init__(self, parent: Kernel):
self.parent = parent
self.type_out = box_spec().to_pyarrow().with_crs(parent._type_in.crs)

def push(self, arr):
parent_result = self.parent.push(arr)
return (
None if parent_result is None else self._old_box_to_new_box(parent_result)
)

def finish(self):
return self._old_box_to_new_box(self.parent.finish())

def _old_box_to_new_box(self, array):
xmin, xmax, ymin, ymax = array.flatten()
storage = pa.StructArray.from_arrays(
[xmin, ymin, xmax, ymax], names=["xmin", "ymin", "xmax", "ymax"]
)
return self.type_out.wrap_array(storage)
Loading
Loading