refactor(geoarrow-pyarrow): Use python type_id and geobuffers implementation

geoarrow-pyarrow/pyproject.toml

@@ -23,7 +23,7 @@ description = ""
 authors = [{name = "Dewey Dunnington", email = "dewey@dunnington.ca"}]
 license = {text = "Apache-2.0"}
 requires-python = ">=3.9"
-dependencies = ["pyarrow >= 14.0.2", "geoarrow-types >= 0.3.0", "geoarrow-c >= 0.3.0"]
+dependencies = ["pyarrow >= 14.0.2", "geoarrow-types >= 0.3.0", "geoarrow-c >= 0.3.1"]
 
 [project.optional-dependencies]
 test = ["pytest", "pandas", "numpy", "geopandas", "pyogrio", "pyproj"]

geoarrow-pyarrow/src/geoarrow/pyarrow/_array.py

@@ -1,4 +1,4 @@
-from geoarrow.pyarrow._kernel import Kernel, _geoarrow_c
+from geoarrow.pyarrow._kernel import Kernel
 from geoarrow.pyarrow._type import (
     GeometryExtensionType,
     large_wkb,
@@ -8,22 +8,12 @@
 )
 
 import pyarrow as pa
+from pyarrow import types as pa_types
 
 
 class GeometryExtensionArray(pa.ExtensionArray):
     def geobuffers(self):
-        import numpy as np
-
-        lib = _geoarrow_c()
-
-        cschema = lib.SchemaHolder()
-        self.type._export_to_c(cschema._addr())
-        carray = lib.ArrayHolder()
-        self._export_to_c(carray._addr())
-
-        array_view = lib.CArrayView(carray, cschema)
-        buffers = array_view.buffers()
-        return [np.array(b) if b is not None else None for b in buffers]
+        return _geobuffers(self.storage)
 
     def __repr__(self):
         # Pretty WKT printing needs geoarrow-c
@@ -161,3 +151,49 @@ def array(obj, type_=None, *args, **kwargs) -> GeometryExtensionArray:
     # Eventually we will be able to handle more types (e.g., parse wkt or wkb
     # into a geoarrow type)
     raise TypeError(f"Can't create geoarrow.array for type {type_}")
+
+
+def _geobuffers(array, include_validity=True):
+    """Collect the physical buffers of ``array`` in depth-first order."""
+    import numpy as np
+
+    # While we may be able to be slightly smarter about not copying for slices
+    # where we can salvage pieces of the original buffer without copying, this
+    # line ensures correctness in the case of a sliced array (where buffers
+    # require taking into account the parent or multiple parents' offsets)
+    if array.offset != 0:
+        array = pa.concat_arrays([array])
+
+    type_ = array.type
+    buffers = array.buffers()
+    out = []
+
+    if include_validity:
+        out.append(_buffer_to_numpy(buffers[0], np.uint8))
+
+    if pa_types.is_struct(type_):
+        for i in range(type_.num_fields):
+            out.extend(_geobuffers(array.field(i), include_validity=False))
+    elif pa_types.is_list(type_) or pa_types.is_large_list(type_):
+        offset_dtype = np.int64 if pa_types.is_large_list(type_) else np.int32
+        out.append(_buffer_to_numpy(buffers[1], offset_dtype))
+        out.extend(_geobuffers(array.values, include_validity=False))
+    elif pa_types.is_fixed_size_list(type_):
+        out.extend(_geobuffers(array.values, include_validity=False))
+    elif pa_types.is_float32(type_):
+        out.append(_buffer_to_numpy(buffers[1], np.float32))
+    elif pa_types.is_float64(type_):
+        out.append(_buffer_to_numpy(buffers[1], np.float64))
+    else:
+        raise TypeError(f"Unsupported GeoArrow storage type for geobuffers(): {type_}")
+
+    return out
+
+
+def _buffer_to_numpy(buffer, dtype):
+    import numpy as np
+
+    if buffer is None:
+        return None
+
+    return np.frombuffer(buffer, dtype=dtype)

geoarrow-pyarrow/src/geoarrow/pyarrow/_compute.py

@@ -1,6 +1,6 @@
 from geoarrow.pyarrow import _type
 from geoarrow.pyarrow._array import array
-from geoarrow.pyarrow._kernel import Kernel, _geoarrow_c
+from geoarrow.pyarrow._kernel import Kernel
 from geoarrow.types import (
     CoordType,
     Dimensions,
@@ -215,7 +215,7 @@ def as_wkb(obj, strict_iso_wkb=False):
     obj = as_geoarrow(obj, _type.wkb())
 
     if check_wkb and strict_iso_wkb and _any_ewkb(obj):
-        return push_all(Kernel.as_geoarrow, obj, args={"type_id": 100001})
+        return push_all(Kernel.as_geoarrow, obj, args={"type_out": _type.wkb()})
     else:
         return obj
 
@@ -262,13 +262,7 @@ def as_geoarrow(obj, type=None, coord_type=None, promote_multi=False):
     if obj.type.spec == type.spec:
         return obj
 
-    lib = _geoarrow_c()
-
-    cschema = lib.SchemaHolder()
-    type._export_to_c(cschema._addr())
-    ctype = lib.CVectorType.FromExtension(cschema)
-
-    return push_all(Kernel.as_geoarrow, obj, args={"type_id": ctype.id})
+    return push_all(Kernel.as_geoarrow, obj, args={"type_out": type})
 
 
 def format_wkt(obj, precision=None, max_element_size_bytes=None):

geoarrow-pyarrow/src/geoarrow/pyarrow/_kernel.py

@@ -1,16 +1,17 @@
 import sys
 
 from geoarrow.pyarrow._type import GeometryExtensionType
-from geoarrow.types import box as box_spec
+from geoarrow.types import CoordType, Dimensions, Encoding, GeometryType
+from geoarrow.types import wkb as wkb_spec
+from geoarrow.types import wkt as wkt_spec
 
 import pyarrow as pa
 
 _lazy_lib = None
-_geoarrow_c_version = None
 
 
 def _geoarrow_c():
-    global _lazy_lib, _geoarrow_c_version
+    global _lazy_lib
     if _lazy_lib is None:
         try:
             import geoarrow.c
@@ -19,10 +20,6 @@ def _geoarrow_c():
             raise ImportError("Requested operation requires geoarrow-c") from e
 
         _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
 
@@ -91,11 +88,11 @@ def visit_void_agg(type_in):
 
     @staticmethod
     def as_wkt(type_in):
-        return Kernel.as_geoarrow(type_in, 100003)
+        return Kernel.as_geoarrow(type_in, wkt_spec().to_pyarrow())
 
     @staticmethod
     def as_wkb(type_in):
-        return Kernel.as_geoarrow(type_in, 100001)
+        return Kernel.as_geoarrow(type_in, wkb_spec().to_pyarrow())
 
     @staticmethod
     def format_wkt(type_in, precision=None, max_element_size_bytes=None):
@@ -107,28 +104,20 @@ def format_wkt(type_in, precision=None, max_element_size_bytes=None):
         )
 
     @staticmethod
-    def as_geoarrow(type_in, type_id):
-        return Kernel("as_geoarrow", type_in, type=int(type_id))
+    def as_geoarrow(type_in, type_out):
+        return Kernel("as_geoarrow", type_in, type=_type_id(type_out))
 
     @staticmethod
     def unique_geometry_types_agg(type_in):
         return Kernel("unique_geometry_types_agg", type_in)
 
     @staticmethod
     def box(type_in):
-        kernel = Kernel("box", type_in)
-        if _geoarrow_c_version <= (0, 1, 3):
-            return BoxKernelCompat(kernel)
-        else:
-            return kernel
+        return Kernel("box", type_in)
 
     @staticmethod
     def box_agg(type_in):
-        kernel = Kernel("box_agg", type_in)
-        if _geoarrow_c_version <= (0, 1, 3):
-            return BoxKernelCompat(kernel)
-        else:
-            return kernel
+        return Kernel("box_agg", type_in)
 
     @staticmethod
     def _pack_options(options):
@@ -149,27 +138,56 @@ def _pack_options(options):
         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 _type_id(type_out):
+    """Previous versions used geoarrow-c's exposed CVectorType for this; however,
+    the CVectorType has been deprecated in favour of geoarrow.types. This keeps
+    geoarrow-pyarrow working against previous and future versions of geoarrow-c
+    while freeing up geoarrow-c to remove code that was used only here."""
+    spec = type_out.spec
 
-    def __init__(self, parent: Kernel):
-        self.parent = parent
-        self.type_out = box_spec().to_pyarrow().with_crs(parent._type_in.crs)
+    if spec.encoding in _SERIALIZED_TYPE_IDS:
+        return _SERIALIZED_TYPE_IDS[spec.encoding]
 
-    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())
+    if spec.encoding != Encoding.GEOARROW:
+        raise ValueError(f"Unsupported GeoArrow encoding: {spec.encoding}")
 
-    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"]
+    try:
+        return (
+            _GEOMETRY_TYPE_IDS[spec.geometry_type]
+            + _DIMENSION_TYPE_ID_OFFSETS[spec.dimensions]
+            + _COORD_TYPE_ID_OFFSETS[spec.coord_type]
         )
-        return self.type_out.wrap_array(storage)
+    except KeyError as e:
+        raise ValueError(f"Unsupported GeoArrow type: {spec}") from e
+
+
+_SERIALIZED_TYPE_IDS = {
+    Encoding.WKB: 100001,
+    Encoding.LARGE_WKB: 100002,
+    Encoding.WKT: 100003,
+    Encoding.LARGE_WKT: 100004,
+    Encoding.WKB_VIEW: 100005,
+    Encoding.WKT_VIEW: 100006,
+}
+
+_GEOMETRY_TYPE_IDS = {
+    GeometryType.POINT: 1,
+    GeometryType.LINESTRING: 2,
+    GeometryType.POLYGON: 3,
+    GeometryType.MULTIPOINT: 4,
+    GeometryType.MULTILINESTRING: 5,
+    GeometryType.MULTIPOLYGON: 6,
+    GeometryType.BOX: 990,
+}
+
+_DIMENSION_TYPE_ID_OFFSETS = {
+    Dimensions.XY: 0,
+    Dimensions.XYZ: 1000,
+    Dimensions.XYM: 2000,
+    Dimensions.XYZM: 3000,
+}
+
+_COORD_TYPE_ID_OFFSETS = {
+    CoordType.SEPARATED: 0,
+    CoordType.INTERLEAVED: 10000,
+}

geoarrow-pyarrow/src/geoarrow/pyarrow/dataset.py

@@ -6,9 +6,9 @@
 
 from concurrent.futures import ThreadPoolExecutor, wait
 
-from geoarrow.c.lib import CoordType
 from geoarrow.pyarrow._kernel import Kernel
 from geoarrow.pyarrow._type import GeometryExtensionType, wkb, wkt
+from geoarrow.types import CoordType
 
 import pyarrow as _pa
 import pyarrow.compute as _compute
@@ -385,7 +385,7 @@ def _wrap_parent(self, filtered_parent, fragment_indices):
 
     def _build_index(self, geometry_columns, geometry_types, num_threads=None):
         can_use_statistics = [
-            type.coord_type == CoordType.SEPARATE for type in self.geometry_types
+            type.coord_type == CoordType.SEPARATED for type in self.geometry_types
         ]
 
         if not self._use_column_statistics or not any(can_use_statistics):
@@ -454,16 +454,18 @@ def _parquet_field_boxes(self, parquet_indices):
                     boxes[i].append(
                         {
                             "xmin": stats_x.min,
-                            "xmax": stats_x.max,
                             "ymin": stats_y.min,
+                            "xmax": stats_x.max,
                             "ymax": stats_y.max,
                         }
                     )
 
             last_row_group = row_group
 
-        type_field_names = ["xmin", "xmax", "ymin", "ymax"]
-        type_fields = [_pa.field(name, _pa.float64()) for name in type_field_names]
+        type_field_names = ["xmin", "ymin", "xmax", "ymax"]
+        type_fields = [
+            _pa.field(name, _pa.float64(), nullable=False) for name in type_field_names
+        ]
         type = _pa.struct(type_fields)
         return [_pa.array(box, type=type) for box in boxes]