feat(geoarrow-types): Add support for box type

geoarrow-types/src/geoarrow/types/__init__.py

@@ -17,6 +17,7 @@
     large_wkb,
     wkt,
     large_wkt,
+    box,
     point,
     linestring,
     polygon,
@@ -42,6 +43,7 @@
     "wkt",
     "large_wkt",
     "geoarrow",
+    "box",
     "point",
     "linestring",
     "polygon",

geoarrow-types/src/geoarrow/types/constants.py

@@ -142,6 +142,9 @@ class GeometryType(TypeSpecEnum):
     GEOMETRYCOLLECTION = 7
     """Geometry collection geometry type"""
 
+    BOX = 990
+    """Box geometry type"""
+
     @classmethod
     def _common2(cls, lhs, rhs):
         out = super()._common2(lhs, rhs)

geoarrow-types/src/geoarrow/types/type_pyarrow.py

@@ -265,6 +265,18 @@ class GeometryCollectionUnionType(GeometryExtensionType):
     _extension_name = "geoarrow.geometrycollection"
 
 
+class BoxType(GeometryExtensionType):
+    """Extension type whose storage is an array of boxes stored
+    as a struct with two children per dimension.
+    """
+
+    _extension_name = "geoarrow.box"
+
+    def from_geobuffers(self, validity, xmin, ymin, *bounds):
+        storage = _from_buffers_box(self.storage_type, validity, xmin, ymin, *bounds)
+        return self.wrap_array(storage)
+
+
 class PointType(GeometryExtensionType):
     """Extension type whose storage is an array of points stored
     as either a struct with one child per dimension or a fixed-size
@@ -468,6 +480,7 @@ def register_extension_types(lazy: bool = True) -> None:
     all_types = [
         type_spec(Encoding.WKT).to_pyarrow(),
         type_spec(Encoding.WKB).to_pyarrow(),
+        type_spec(Encoding.GEOARROW, GeometryType.BOX).to_pyarrow(),
         type_spec(Encoding.GEOARROW, GeometryType.GEOMETRY).to_pyarrow(),
         type_spec(Encoding.GEOARROW, GeometryType.POINT).to_pyarrow(),
         type_spec(Encoding.GEOARROW, GeometryType.LINESTRING).to_pyarrow(),
@@ -504,12 +517,15 @@ def unregister_extension_types(lazy=True):
     all_type_names = [
         "geoarrow.wkb",
         "geoarrow.wkt",
+        "geoarrow.box",
+        "geoarrow.geometry",
         "geoarrow.point",
         "geoarrow.linestring",
         "geoarrow.polygon",
         "geoarrow.multipoint",
         "geoarrow.multilinestring",
         "geoarrow.multipolygon",
+        "geoarrow.geometrycollection",
     ]
 
     n_unregistered = 0
@@ -624,7 +640,15 @@ def _deserialize_storage(storage_type, extension_name=None, extension_metadata=N
         names, n_dims, parsed_children = params
         n_dims_infer = n_dims
 
-    if names in _DIMS_FROM_NAMES:
+    # Make sure we catch box field names (e.g., xmin, ymin, ...)
+    if names in _BOX_DIMS_FROM_NAMES:
+        if spec.geometry_type != GeometryType.POINT:
+            raise ValueError(
+                f"Expected box names {names} in root type but got nested list"
+            )
+        spec = spec.override(geometry_type=GeometryType.BOX)
+        dims = _BOX_DIMS_FROM_NAMES[names]
+    elif names in _DIMS_FROM_NAMES:
         dims = _DIMS_FROM_NAMES[names]
         if n_dims != dims.count():
             raise ValueError(f"Expected {n_dims} dimensions but got Dimensions.{dims}")
@@ -694,6 +718,13 @@ def _pybuffer_offset(x):
     return len(mv), pa.py_buffer(mv)
 
 
+def _from_buffers_box(type_, validity, *bounds):
+    length = len(bounds[0])
+    validity = pa.py_buffer(validity) if validity is not None else None
+    children = [_from_buffer_ordinate(bound) for bound in bounds]
+    return pa.Array.from_buffers(type_, length, buffers=[validity], children=children)
+
+
 def _from_buffers_point(type_, validity, x, y=None, z_or_m=None, m=None):
     validity = pa.py_buffer(validity) if validity is not None else None
     children = [_from_buffer_ordinate(x)]
@@ -785,6 +816,13 @@ def _from_buffers_multipolygon(
     GeometryType.MULTILINESTRING,
     GeometryType.MULTIPOLYGON,
 ]
+_BOX_DIMS_FROM_NAMES = {
+    ("xmin", "ymin", "xmax", "ymax"): Dimensions.XY,
+    ("xmin", "ymin", "zmin", "xmax", "ymax", "zmax"): Dimensions.XYZ,
+    ("xmin", "ymin", "mmin", "xmax", "ymax", "mmax"): Dimensions.XYM,
+    ("xmin", "ymin", "zmin", "mmin", "xmax", "ymax", "zmax", "mmax"): Dimensions.XYZM,
+}
+_BOX_NAMES_FROM_DIMS = {v: k for k, v in _BOX_DIMS_FROM_NAMES.items()}
 
 
 def _generate_storage_types():
@@ -818,6 +856,13 @@ def _generate_storage_types():
                 storage_type = _nested_type(coord, names)
                 all_storage_types[key] = storage_type
 
+    for dimensions in ALL_DIMENSIONS:
+        storage_type = _nested_type(
+            _struct_fields(_BOX_NAMES_FROM_DIMS[dimensions]), []
+        )
+        key = GeometryType.BOX, CoordType.SEPARATED, dimensions
+        all_storage_types[key] = storage_type
+
     return all_storage_types
 
 
@@ -934,8 +979,9 @@ def _spec_short_repr(spec, ext_name):
 _EXTENSION_CLASSES = {
     "geoarrow.wkb": WkbType,
     "geoarrow.wkt": WktType,
-    "geoarrow.point": PointType,
+    "geoarrow.box": BoxType,
     "geoarrow.geometry": GeometryUnionType,
+    "geoarrow.point": PointType,
     "geoarrow.linestring": LinestringType,
     "geoarrow.polygon": PolygonType,
     "geoarrow.multipoint": MultiPointType,

geoarrow-types/src/geoarrow/types/type_spec.py

@@ -428,6 +428,29 @@ def point(
     )
 
 
+def box(
+    *,
+    dimensions=None,
+    coord_type=None,
+    edge_type=None,
+    crs=crs.UNSPECIFIED,
+) -> TypeSpec:
+    """GeoArrow box
+
+    Create a :class:`TypeSpec` denoting a preference for GeoArrow Box
+    type without an explicit request for dimensions or coordinate type.
+    See :func:`type_spec` for parameter definitions.
+    """
+    return type_spec(
+        encoding=Encoding.GEOARROW,
+        geometry_type=GeometryType.BOX,
+        dimensions=dimensions,
+        coord_type=coord_type,
+        edge_type=edge_type,
+        crs=crs,
+    )
+
+
 def linestring(
     *,
     dimensions=None,
@@ -599,6 +622,7 @@ def type_spec(
 }
 
 _GEOARROW_EXT_NAMES = {
+    GeometryType.BOX: "geoarrow.box",
     GeometryType.GEOMETRY: "geoarrow.geometry",
     GeometryType.POINT: "geoarrow.point",
     GeometryType.LINESTRING: "geoarrow.linestring",

geoarrow-types/tests/test_type_pyarrow.py

@@ -301,6 +301,27 @@ def test_geometry_collection_union_type():
     assert geometry.geometry_type == gt.GeometryType.GEOMETRYCOLLECTION
 
 
+def test_box_array_from_geobuffers():
+    pa_type = gt.box(dimensions=gt.Dimensions.XY).to_pyarrow()
+    arr = pa_type.from_geobuffers(
+        b"\xff",
+        np.array([1.0, 2.0, 3.0]),
+        np.array([4.0, 5.0, 6.0]),
+        np.array([7.0, 8.0, 9.0]),
+        np.array([10.0, 11.0, 12.0]),
+    )
+    assert len(arr) == 3
+    assert arr.type == pa_type
+    assert arr.storage == pa.array(
+        [
+            {"xmin": 1.0, "ymin": 4.0, "xmax": 7.0, "ymax": 10.0},
+            {"xmin": 2.0, "ymin": 5.0, "xmax": 8.0, "ymax": 11.0},
+            {"xmin": 3.0, "ymin": 6.0, "xmax": 9.0, "ymax": 12.0},
+        ],
+        pa_type.storage_type,
+    )
+
+
 def test_point_array_from_geobuffers():
     pa_type = gt.point(dimensions=gt.Dimensions.XYZM).to_pyarrow()
     arr = pa_type.from_geobuffers(
@@ -417,6 +438,7 @@ def test_multipolygon_array_from_geobuffers():
         gt.wkb(),
         gt.large_wkb(),
         # Geometry types
+        gt.box(),
         gt.point(),
         gt.linestring(),
         gt.polygon(),
@@ -433,6 +455,11 @@ def test_multipolygon_array_from_geobuffers():
         gt.point(dimensions="xyz", coord_type="interleaved"),
         gt.point(dimensions="xym", coord_type="interleaved"),
         gt.point(dimensions="xyzm", coord_type="interleaved"),
+        # Box with all dimensions
+        gt.box(dimensions="xy"),
+        gt.box(dimensions="xyz"),
+        gt.box(dimensions="xym"),
+        gt.box(dimensions="xyzm"),
         # Union types
         gt.type_spec(gt.Encoding.GEOARROW, gt.GeometryType.GEOMETRY),
         gt.type_spec(gt.Encoding.GEOARROW, gt.GeometryType.GEOMETRYCOLLECTION),

geoarrow-types/tests/test_type_spec.py

@@ -317,6 +317,9 @@ def test_type_spec_shortcuts():
     assert gt.large_wkt() == TypeSpec(encoding=Encoding.LARGE_WKT)
 
     assert gt.geoarrow() == TypeSpec(encoding=Encoding.GEOARROW)
+    assert gt.box() == TypeSpec(
+        encoding=Encoding.GEOARROW, geometry_type=GeometryType.BOX
+    )
     assert gt.point() == TypeSpec(
         encoding=Encoding.GEOARROW, geometry_type=GeometryType.POINT
     )