Support pyarrow.Table with geoarrow.pyarrow extension types as geometry columns (#218)

Work-in-progress, and haven't yet tested it.

The idea is that when encountering a `geoarrow.pyarrow` extension type,
to convert this to its raw storage field with metadata (i.e. how the
extension type looks like in IPC or when the python wrapper class is not
registered). Short term, this is the easiest approach, because a lot of
the existing code relies on checking eg the "ARROW:extension:name" in
the field metadata.

---------

Co-authored-by: Kyle Barron <[email protected]>

Author: jorisvandenbossche

lonboard/_geoarrow/sanitize.py

@@ -0,0 +1,59 @@
+"""Remove custom geoarrow.pyarrow types from input geoarrow data
+"""
+import json
+from typing import Tuple
+
+import pyarrow as pa
+from pyproj import CRS
+
+
+def sanitize_table(table: pa.Table) -> pa.Table:
+    """
+    Convert any registered geoarrow.pyarrow extension fields and arrays to plain
+    metadata
+    """
+    for field_idx in range(len(table.schema)):
+        field = table.field(field_idx)
+        column = table.column(field_idx)
+
+        if isinstance(field.type, pa.ExtensionType):
+            assert all(isinstance(chunk, pa.ExtensionArray) for chunk in column.chunks)
+            new_field, new_column = sanitize_column(field, column)
+            table = table.set_column(field_idx, new_field, new_column)
+
+    return table
+
+
+def sanitize_column(
+    field: pa.Field, column: pa.ChunkedArray
+) -> Tuple[pa.Field, pa.ChunkedArray]:
+    """
+    Convert a registered geoarrow.pyarrow extension field and column to plain metadata
+    """
+    import geoarrow.pyarrow as gap
+
+    extension_metadata = {}
+    if field.type.crs:
+        extension_metadata["crs"] = CRS.from_user_input(field.type.crs).to_json()
+
+    if field.type.edge_type == gap.EdgeType.SPHERICAL:
+        extension_metadata["edges"] = "spherical"
+
+    metadata = {
+        "ARROW:extension:name": field.type.extension_name,
+    }
+    if extension_metadata:
+        metadata["ARROW:extension:metadata"] = json.dumps(extension_metadata)
+
+    new_field = pa.field(
+        field.name, field.type.storage_type, nullable=field.nullable, metadata=metadata
+    )
+
+    new_chunks = []
+    for chunk in column.chunks:
+        if hasattr(chunk, "storage"):
+            new_chunks.append(chunk.storage)
+        else:
+            new_chunks.append(chunk.cast(new_field.type))
+
+    return new_field, pa.chunked_array(new_chunks)

lonboard/_layer.py

@@ -23,6 +23,7 @@
 from lonboard._geoarrow.ops import reproject_table
 from lonboard._geoarrow.ops.bbox import Bbox, total_bounds
 from lonboard._geoarrow.ops.centroid import WeightedCentroid, weighted_centroid
+from lonboard._geoarrow.sanitize import sanitize_table
 from lonboard._serialization import infer_rows_per_chunk
 from lonboard._utils import auto_downcast as _auto_downcast
 from lonboard._utils import get_geometry_column_index, remove_extension_kwargs
@@ -231,6 +232,13 @@ class BaseArrowLayer(BaseLayer):
     def __init__(
         self, *, table: pa.Table, _rows_per_chunk: Optional[int] = None, **kwargs
     ):
+        # Check for Arrow PyCapsule Interface
+        # https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
+        if not isinstance(table, pa.Table) and hasattr(table, "__arrow_c_stream__"):
+            table = pa.table(table)
+
+        table = sanitize_table(table)
+
         # Reproject table to WGS84 if needed
         # Note this must happen before calculating the default viewport
         table = reproject_table(table, to_crs=OGC_84)

lonboard/traits.py

@@ -21,6 +21,7 @@
     ACCESSOR_SERIALIZATION,
     TABLE_SERIALIZATION,
 )
+from lonboard._utils import get_geometry_column_index
 
 
 # This is a custom subclass of traitlets.TraitType because its `error` method ignores
@@ -139,21 +140,19 @@ def __init__(
         )
 
     def validate(self, obj: Self, value: Any):
-        # Check for Arrow PyCapsule Interface
-        # https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
-        if not isinstance(value, pa.Table) and hasattr(value, "__arrow_c_stream__"):
-            value = pa.table(value)
-
         if not isinstance(value, pa.Table):
             self.error(obj, value)
 
         allowed_geometry_types = self.metadata.get("allowed_geometry_types")
+        # No restriction on the allowed geometry types in this table
         if not allowed_geometry_types:
             return value
 
-        geometry_extension_type = value.schema.field("geometry").metadata.get(
+        geom_col_idx = get_geometry_column_index(value.schema)
+        geometry_extension_type = value.schema.field(geom_col_idx).metadata.get(
             b"ARROW:extension:name"
         )
+
         if (
             allowed_geometry_types
             and geometry_extension_type not in allowed_geometry_types

poetry.lock

@@ -32,6 +32,7 @@ black = "^23.10.1"
 geoarrow-rust-core = "^0.1.0"
 geodatasets = "^2023.12.0"
 pyogrio = "^0.7.2"
+geoarrow-pyarrow = "^0.1.1"
 
 [tool.poetry.group.docs.dependencies]
 mkdocs = "^1.4.3"

pyproject.toml

@@ -1,5 +1,6 @@
 import geopandas as gpd
 import numpy as np
+import pyarrow as pa
 import pytest
 import shapely
 from traitlets import TraitError
@@ -58,3 +59,16 @@ def test_layer_outside_4326_range():
 
     with pytest.raises(ValueError, match="outside of WGS84 bounds"):
         _layer = ScatterplotLayer.from_geopandas(gdf)
+
+
+def test_layer_from_geoarrow_pyarrow():
+    ga = pytest.importorskip("geoarrow.pyarrow")
+
+    points = gpd.GeoSeries(shapely.points([1, 2], [3, 4]))
+
+    # convert to geoarrow.pyarrow Table (currently requires to ensure interleaved
+    # coordinates manually)
+    points = ga.with_coord_type(ga.as_geoarrow(points), ga.CoordType.INTERLEAVED)
+    table = pa.table({"geometry": points})
+
+    _layer = ScatterplotLayer(table=table)