diff --git a/geoarrow-types/src/geoarrow/types/__init__.py b/geoarrow-types/src/geoarrow/types/__init__.py index 029b17a..4dcd822 100644 --- a/geoarrow-types/src/geoarrow/types/__init__.py +++ b/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", diff --git a/geoarrow-types/src/geoarrow/types/constants.py b/geoarrow-types/src/geoarrow/types/constants.py index 8866780..40bb01b 100644 --- a/geoarrow-types/src/geoarrow/types/constants.py +++ b/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) diff --git a/geoarrow-types/src/geoarrow/types/type_pyarrow.py b/geoarrow-types/src/geoarrow/types/type_pyarrow.py index 9cdcfb6..eea929d 100644 --- a/geoarrow-types/src/geoarrow/types/type_pyarrow.py +++ b/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, diff --git a/geoarrow-types/src/geoarrow/types/type_spec.py b/geoarrow-types/src/geoarrow/types/type_spec.py index ee07afe..3d07ce8 100644 --- a/geoarrow-types/src/geoarrow/types/type_spec.py +++ b/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", diff --git a/geoarrow-types/tests/test_type_pyarrow.py b/geoarrow-types/tests/test_type_pyarrow.py index 8c0da33..f92b2e6 100644 --- a/geoarrow-types/tests/test_type_pyarrow.py +++ b/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), diff --git a/geoarrow-types/tests/test_type_spec.py b/geoarrow-types/tests/test_type_spec.py index d6cd7d5..56d8b82 100644 --- a/geoarrow-types/tests/test_type_spec.py +++ b/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 )