chore: Ensure support for GeoArrow 0.2 CRS (#797)

Closes #791

Author: kylebarron

lonboard/_geoarrow/crs.py

@@ -1,18 +1,56 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+
+from pyproj import CRS
 
 if TYPE_CHECKING:
     from arro3.core import Field
 
 
 # Note: According to the spec, if the metadata key exists, its value should never be
 # `null` or an empty dict, but we still check for those to be safe
-def get_field_crs(field: Field) -> dict | None:
-    extension_metadata_value = field.metadata.get(b"ARROW:extension:metadata")
+def get_field_crs(field: Field) -> CRS | None:
+    extension_metadata_value = field.metadata_str.get("ARROW:extension:metadata")
     if not extension_metadata_value:
         return None
 
     extension_metadata = json.loads(extension_metadata_value)
-    return extension_metadata.get("crs")
+    return parse_metadata(extension_metadata)
+
+
+def parse_metadata(extension_metadata: dict[str, Any]) -> CRS | None:
+    crs_val = extension_metadata.get("crs")
+    crs_type = extension_metadata.get("crs_type")
+
+    if crs_type == "projjson":
+        assert crs_val is not None, "CRS value must be provided for projjson type"
+        return CRS.from_json_dict(crs_val)
+
+    if crs_type == "wkt2:2019":
+        assert crs_val is not None, "CRS value must be provided for WKT2:2019 type"
+        return CRS.from_wkt(crs_val)
+
+    if crs_type == "authority_code":
+        assert crs_val is not None, "CRS value must be provided for authority code type"
+        assert ":" in crs_val, "Authority code must be in the format 'authority:code'"
+        assert isinstance(crs_val, str), "CRS value must be a string"
+        return CRS.from_authority(*crs_val.split(":", 1))
+
+    return CRS.from_user_input(crs_val) if crs_val is not None else None
+
+
+def serialize_crs(crs: CRS | None) -> dict[str, Any] | None:
+    """Serialize a CRS to GeoArrow metadata.
+
+    The returned GeoArrow metadata should be JSON-serialized within the
+    ARROW:extension:metadata key.
+    """
+    if crs is None:
+        return None
+
+    return {
+        "crs": crs.to_json_dict(),
+        "crs_type": "projjson",
+    }

lonboard/_geoarrow/extension_types.py

@@ -7,10 +7,13 @@
 import numpy as np
 from arro3.core import Array, DataType, Field, fixed_size_list_array, list_array
 
+from lonboard._geoarrow.crs import serialize_crs
+
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
     from numpy.typing import NDArray
+    from pyproj import CRS
 
 
 class CoordinateDimension(str, Enum):
@@ -200,7 +203,7 @@ def construct_geometry_array(  # noqa: PLR0915
     include_z: bool | None = None,  # noqa: FBT001
     *,
     field_name: str = "geometry",
-    crs: dict | None = None,
+    crs: CRS | None = None,
 ) -> tuple[Field, Array]:
     import shapely
     from shapely import GeometryType
@@ -220,7 +223,7 @@ def construct_geometry_array(  # noqa: PLR0915
 
     extension_metadata: dict[str, str] = {}
     if crs is not None:
-        extension_metadata["ARROW:extension:metadata"] = json.dumps({"crs": crs})
+        extension_metadata["ARROW:extension:metadata"] = json.dumps(serialize_crs(crs))
 
     if geom_type == GeometryType.POINT:
         arrow_coords = fixed_size_list_array(coords.ravel("C"), len(dims)).cast(

lonboard/_geoarrow/geopandas_interop.py

@@ -32,7 +32,7 @@ def geopandas_to_geoarrow(
     pyarrow_table = pa.Table.from_pandas(df_attr, preserve_index=preserve_index)
     field, geom_arr = construct_geometry_array(
         np.array(gdf.geometry),
-        crs=gdf.crs.to_json_dict() if gdf.crs is not None else None,
+        crs=gdf.crs,
     )
     return Table.from_arrow(pyarrow_table).append_column(
         field,

lonboard/_geoarrow/ops/reproject.py

@@ -96,13 +96,11 @@ def reproject_column(
 
     """
     extension_type_name = field.metadata[b"ARROW:extension:name"]
-    crs_str = get_field_crs(field)
-    if crs_str is None:
+    existing_crs = get_field_crs(field)
+    if existing_crs is None:
         no_crs_warning()
         return field, column
 
-    existing_crs = CRS(crs_str)
-
     if existing_crs == to_crs:
         return field, column
 

tests/geoarrow/__init__.py

@@ -0,0 +1,124 @@
+import json
+
+from arro3.core import Field
+from pyproj import CRS
+
+from lonboard._geoarrow.crs import get_field_crs
+from lonboard._geoarrow.extension_types import CoordinateDimension, coord_storage_type
+
+
+def parse_geoarrow_projjson_crs():
+    expected = CRS.from_epsg(4326)
+    meta = {
+        "crs": expected.to_json_dict(),
+        "crs_type": "projjson",
+    }
+    geoarrow_meta = {
+        "ARROW:extension:name": "geoarrow.point",
+        "ARROW:extension:metadata": json.dumps(meta),
+    }
+    field = Field(
+        name="geometry",
+        type=coord_storage_type(interleaved=True, dims=CoordinateDimension.XY),
+        nullable=True,
+        metadata=geoarrow_meta,
+    )
+    assert get_field_crs(field) == expected
+
+
+def parse_geoarrow_wkt_crs():
+    expected = CRS.from_epsg(4326)
+    meta = {
+        "crs": expected.to_wkt(version="WKT2:2019"),
+        "crs_type": "wkt2:2019",
+    }
+    geoarrow_meta = {
+        "ARROW:extension:name": "geoarrow.point",
+        "ARROW:extension:metadata": json.dumps(meta),
+    }
+    field = Field(
+        name="geometry",
+        type=coord_storage_type(interleaved=True, dims=CoordinateDimension.XY),
+        nullable=True,
+        metadata=geoarrow_meta,
+    )
+    assert get_field_crs(field) == expected
+
+
+def parse_geoarrow_authority_crs():
+    expected = CRS.from_epsg(4326)
+    meta = {
+        "crs": ":".join(expected.to_authority()),
+        "crs_type": "authority_code",
+    }
+    geoarrow_meta = {
+        "ARROW:extension:name": "geoarrow.point",
+        "ARROW:extension:metadata": json.dumps(meta),
+    }
+    field = Field(
+        name="geometry",
+        type=coord_storage_type(interleaved=True, dims=CoordinateDimension.XY),
+        nullable=True,
+        metadata=geoarrow_meta,
+    )
+    assert get_field_crs(field) == expected
+
+
+def parse_geoarrow_srid_crs():
+    expected = CRS.from_epsg(4326)
+    meta = {
+        "crs": str(expected.to_epsg()),
+        "crs_type": "srid",
+    }
+    geoarrow_meta = {
+        "ARROW:extension:name": "geoarrow.point",
+        "ARROW:extension:metadata": json.dumps(meta),
+    }
+    field = Field(
+        name="geometry",
+        type=coord_storage_type(interleaved=True, dims=CoordinateDimension.XY),
+        nullable=True,
+        metadata=geoarrow_meta,
+    )
+    assert get_field_crs(field) == expected
+
+
+def parse_geoarrow_unknown_crs_type():
+    expected = CRS.from_epsg(4326)
+    meta = {
+        "crs": expected.to_wkt(),
+    }
+    geoarrow_meta = {
+        "ARROW:extension:name": "geoarrow.point",
+        "ARROW:extension:metadata": json.dumps(meta),
+    }
+    field = Field(
+        name="geometry",
+        type=coord_storage_type(interleaved=True, dims=CoordinateDimension.XY),
+        nullable=True,
+        metadata=geoarrow_meta,
+    )
+    assert get_field_crs(field) == expected
+
+
+def parse_geoarrow_no_crs():
+    geoarrow_meta = {
+        "ARROW:extension:name": "geoarrow.point",
+        "ARROW:extension:metadata": json.dumps({}),
+    }
+    field = Field(
+        name="geometry",
+        type=coord_storage_type(interleaved=True, dims=CoordinateDimension.XY),
+        nullable=True,
+        metadata=geoarrow_meta,
+    )
+    assert get_field_crs(field) is None
+
+    geoarrow_meta = {"ARROW:extension:name": "geoarrow.point"}
+    field = Field(
+        name="geometry",
+        type=coord_storage_type(interleaved=True, dims=CoordinateDimension.XY),
+        nullable=True,
+        metadata=geoarrow_meta,
+    )
+    assert get_field_crs(field) is None