feat: Add A5Layer (#1001)

@felixpalmer 

Closes #968

<img width="529" height="616" alt="image"
src="https://github.com/user-attachments/assets/abc29e32-4777-4262-a290-a68c57e43373"
/>

Author: kylebarron

lonboard/__init__.py

@@ -5,6 +5,7 @@
 from ._version import __version__
 from ._viz import viz
 from .layer import (
+    A5Layer,
     ArcLayer,
     BaseArrowLayer,
     BaseLayer,
@@ -21,6 +22,7 @@
 )
 
 __all__ = [
+    "A5Layer",
     "ArcLayer",
     "BaseArrowLayer",
     "BaseLayer",

lonboard/layer/__init__.py

@@ -7,6 +7,7 @@
 #   Then the default value in the JS GeoArrow layer (defined in
 #   `@geoarrow/deck.gl-layers`) will be used.
 
+from ._a5 import A5Layer
 from ._arc import ArcLayer
 from ._base import BaseArrowLayer, BaseLayer
 from ._bitmap import BitmapLayer, BitmapTileLayer
@@ -20,6 +21,7 @@
 from ._trips import TripsLayer
 
 __all__ = [
+    "A5Layer",
     "ArcLayer",
     "BaseArrowLayer",
     "BaseLayer",

lonboard/layer/_a5.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import traitlets as t
+
+from lonboard._utils import auto_downcast as _auto_downcast
+
+# Important to import from ._polygon to avoid circular imports
+from lonboard.layer._polygon import PolygonLayer
+from lonboard.traits import A5Accessor, ArrowTableTrait
+
+if TYPE_CHECKING:
+    import sys
+
+    import pandas as pd
+    from arro3.core.types import ArrowStreamExportable
+
+    from lonboard.types.layer import A5AccessorInput, A5LayerKwargs
+
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
+    if sys.version_info >= (3, 12):
+        from typing import Unpack
+    else:
+        from typing_extensions import Unpack
+
+
+class A5Layer(PolygonLayer):
+    """The `A5Layer` renders filled and/or stroked polygons based on the [A5](https://a5geo.org) geospatial indexing system."""
+
+    def __init__(
+        self,
+        table: ArrowStreamExportable,
+        *,
+        get_pentagon: A5AccessorInput,
+        _rows_per_chunk: int | None = None,
+        **kwargs: Unpack[A5LayerKwargs],
+    ) -> None:
+        """Create a new A5Layer.
+
+        Args:
+            table: An Arrow table with properties to associate with the A5 pentagons.
+
+        Keyword Args:
+            get_pentagon: The cell identifier of each A5 pentagon.
+            kwargs: Extra args passed down as A5Layer attributes.
+
+        """
+        super().__init__(
+            table=table,
+            get_pentagon=get_pentagon,
+            _rows_per_chunk=_rows_per_chunk,
+            **kwargs,
+        )
+
+    @classmethod
+    def from_pandas(
+        cls,
+        df: pd.DataFrame,
+        *,
+        get_pentagon: A5AccessorInput,
+        auto_downcast: bool = True,
+        **kwargs: Unpack[A5LayerKwargs],
+    ) -> Self:
+        """Create a new A5Layer from a pandas DataFrame.
+
+        Args:
+            df: a Pandas DataFrame with properties to associate with A5 pentagons.
+
+        Keyword Args:
+            get_pentagon: A5 cell identifier of each A5 hexagon.
+            auto_downcast: Whether to save memory on input by casting to smaller types. Defaults to True.
+            kwargs: Extra args passed down as A5Layer attributes.
+
+        """
+        try:
+            import pyarrow as pa
+        except ImportError as e:
+            raise ImportError(
+                "pyarrow required for converting GeoPandas to arrow.\n"
+                "Run `pip install pyarrow`.",
+            ) from e
+
+        if auto_downcast:
+            # Note: we don't deep copy because we don't need to clone geometries
+            df = _auto_downcast(df.copy())  # type: ignore
+
+        table = pa.Table.from_pandas(df)
+        return cls(table, get_pentagon=get_pentagon, **kwargs)
+
+    _layer_type = t.Unicode("a5").tag(sync=True)
+
+    table = ArrowTableTrait(geometry_required=False)
+    """An Arrow table with properties to associate with the A5 pentagons.
+
+    If you have a Pandas `DataFrame`, use
+    [`from_pandas`][lonboard.A5Layer.from_pandas] instead.
+    """
+
+    get_pentagon = A5Accessor()
+    """The cell identifier of each A5 pentagon.
+
+    Accepts either an array of strings or uint64 integers representing A5 cell IDs.
+
+    - Type: [A5Accessor][lonboard.traits.A5Accessor]
+    """

lonboard/layer/_h3.py

@@ -130,10 +130,10 @@ def __init__(
         """Create a new H3HexagonLayer.
 
         Args:
-            table: _description_
+            table: An Arrow table with properties to associate with the H3 hexagons.
 
         Keyword Args:
-            get_hexagon: _description_
+            get_hexagon: The cell identifier of each H3 hexagon.
             kwargs: Extra args passed down as H3HexagonLayer attributes.
 
         """
@@ -170,12 +170,6 @@ def from_pandas(
             auto_downcast: Whether to save memory on input by casting to smaller types. Defaults to True.
             kwargs: Extra args passed down as H3HexagonLayer attributes.
 
-        Raises:
-            ImportError: _description_
-
-        Returns:
-            _description_
-
         """
         try:
             import pyarrow as pa

lonboard/traits/__init__.py

@@ -4,6 +4,7 @@
 documentation on how to define new traitlet types.
 """
 
+from ._a5 import A5Accessor
 from ._base import FixedErrorTraitType, VariableLengthTuple
 from ._color import ColorAccessor
 from ._extensions import DashArrayAccessor, FilterValueAccessor
@@ -17,6 +18,7 @@
 from ._timestamp import TimestampAccessor
 
 __all__ = [
+    "A5Accessor",
     "ArrowTableTrait",
     "BasemapUrl",
     "ColorAccessor",

lonboard/traits/_a5.py

@@ -0,0 +1,128 @@
+# ruff: noqa: SLF001
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from arro3.core import Array, ChunkedArray, DataType
+
+from lonboard._h3._str_to_h3 import str_to_h3
+from lonboard._serialization import ACCESSOR_SERIALIZATION
+from lonboard.traits._base import FixedErrorTraitType
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from numpy.typing import NDArray
+    from traitlets.traitlets import TraitType
+
+    from lonboard.layer import BaseArrowLayer
+
+
+class A5Accessor(FixedErrorTraitType):
+    """A trait to validate A5 cell input.
+
+    Various input is allowed:
+
+    - A numpy `ndarray` with an object, S16, or uint64 data type.
+    - A pandas `Series` with an object or uint64 data type.
+    - A pyarrow string, large string, string view array, or uint64 array, or a chunked array of those types.
+    - Any Arrow string, large string, string view array, or uint64 array, or a chunked array of those types from a library that implements the [Arrow PyCapsule
+      Interface](https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html).
+    """
+
+    default_value = None
+    info_text = (
+        "a float value or numpy ndarray or Arrow array representing an array of floats"
+    )
+
+    def __init__(
+        self: TraitType,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self.tag(sync=True, **ACCESSOR_SERIALIZATION)
+
+    def _pandas_to_numpy(
+        self,
+        obj: BaseArrowLayer,
+        value: pd.Series,
+    ) -> NDArray[np.str_] | NDArray[np.uint64]:
+        """Cast pandas Series to numpy ndarray."""
+        if isinstance(value.dtype, np.dtype) and np.issubdtype(value.dtype, np.integer):
+            return np.asarray(value, dtype=np.uint64)
+
+        if not isinstance(value.dtype, np.dtype) or not np.issubdtype(
+            value.dtype,
+            np.object_,
+        ):
+            self.error(
+                obj,
+                value,
+                info="A5 Pandas series not object or uint64 dtype.",
+            )
+
+        if not (value.str.len() == 16).all():
+            self.error(
+                obj,
+                value,
+                info="A5 Pandas series not all 16 characters long.",
+            )
+
+        return np.asarray(value, dtype="S16")
+
+    def _numpy_to_arrow(self, obj: BaseArrowLayer, value: np.ndarray) -> ChunkedArray:
+        if np.issubdtype(value.dtype, np.uint64):
+            return ChunkedArray([value])
+
+        if np.issubdtype(value.dtype, np.object_):
+            if {len(v) for v in value} != {16}:
+                self.error(
+                    obj,
+                    value,
+                    info="numpy object array not all 16 characters long",
+                )
+
+            value = np.asarray(value, dtype="S16")
+
+        if not np.issubdtype(value.dtype, np.dtype("S16")):
+            self.error(obj, value, info="numpy array not object, str, or uint64 dtype")
+
+        a5_uint8_array = str_to_h3(value)
+        return ChunkedArray([a5_uint8_array])
+
+    def validate(self, obj: BaseArrowLayer, value: Any) -> ChunkedArray:
+        # pandas Series
+        if (
+            value.__class__.__module__.startswith("pandas")
+            and value.__class__.__name__ == "Series"
+        ):
+            value = self._pandas_to_numpy(obj, value)
+
+        if isinstance(value, np.ndarray):
+            value = self._numpy_to_arrow(obj, value)
+        elif hasattr(value, "__arrow_c_array__"):
+            value = ChunkedArray([Array.from_arrow(value)])
+        elif hasattr(value, "__arrow_c_stream__"):
+            value = ChunkedArray.from_arrow(value)
+        else:
+            self.error(obj, value)
+
+        assert isinstance(value, ChunkedArray)
+
+        if (
+            DataType.is_string(value.type)
+            or DataType.is_large_string(value.type)
+            or DataType.is_string_view(value.type)
+        ):
+            value = self._numpy_to_arrow(obj, value.to_numpy())
+
+        if not DataType.is_uint64(value.type):
+            self.error(
+                obj,
+                value,
+                info="A5 Arrow array must be uint64 type.",
+            )
+
+        return value.rechunk(max_chunksize=obj._rows_per_chunk)

lonboard/types/layer.py

@@ -62,6 +62,7 @@
     ArrowArrayExportable,
     ArrowStreamExportable,
 ]
+A5AccessorInput = H3AccessorInput
 NormalAccessorInput = Union[
     list[int],
     tuple[int, int, int],
@@ -169,25 +170,6 @@ class ColumnLayerKwargs(BaseLayerKwargs, total=False):
     get_line_width: FloatAccessorInput
 
 
-class H3HexagonLayerKwargs(BaseLayerKwargs, total=False):
-    high_precision: bool
-    stroked: bool
-    filled: bool
-    extruded: bool
-    wireframe: bool
-    elevation_scale: IntFloat
-    line_width_units: Units
-    line_width_scale: IntFloat
-    line_width_min_pixels: IntFloat
-    line_width_max_pixels: IntFloat
-    line_joint_rounded: bool
-    line_miter_limit: IntFloat
-    get_fill_color: ColorAccessorInput
-    get_line_color: ColorAccessorInput
-    get_line_width: FloatAccessorInput
-    get_elevation: FloatAccessorInput
-
-
 class PathLayerKwargs(BaseLayerKwargs, total=False):
     width_units: Units
     width_scale: IntFloat
@@ -226,6 +208,14 @@ class PolygonLayerKwargs(BaseLayerKwargs, total=False):
     get_elevation: FloatAccessorInput
 
 
+class H3HexagonLayerKwargs(PolygonLayerKwargs, total=False):
+    high_precision: bool
+
+
+class A5LayerKwargs(PolygonLayerKwargs, total=False):
+    pass
+
+
 class ScatterplotLayerKwargs(BaseLayerKwargs, total=False):
     radius_units: Units
     radius_scale: IntFloat