arro3 followups (#596)

kylebarron · web-flow · commit 6b4073b4f637 · 2024-08-14T19:17:24.000Z
### Change list - Rename `PyarrowTableTrait` to `ArrowTableTrait` - Clearer `ImportError` exception messages for code paths that still use pyarrow - Use Arrow exportable types from arro3, because those types are [better documented online](https://kylebarron.dev/arro3/latest/api/core/types/) - Add arro3 for docs hyperlinking
diff --git a/docs/api/traits.md b/docs/api/traits.md
@@ -6,4 +6,4 @@
 ::: lonboard.traits.FilterValueAccessor
 ::: lonboard.traits.NormalAccessor
 ::: lonboard.traits.PointAccessor
-::: lonboard.traits.PyarrowTableTrait
+::: lonboard.traits.ArrowTableTrait
diff --git a/lonboard/_cli.py b/lonboard/_cli.py
@@ -5,7 +5,6 @@
 from typing import Dict, List, Optional
 
 import click
-import pyarrow.parquet as pq
 from arro3.core import Table
 from pyproj import CRS
 
@@ -65,6 +64,14 @@ def read_geoparquet(path: Path) -> Table:
     Args:
         path: Path to GeoParquet file
     """
+    try:
+        import pyarrow.parquet as pq
+    except ImportError as e:
+        raise ImportError(
+            "pyarrow currently required for reading GeoParquet files.\n"
+            "Run `pip install pyarrow`."
+        ) from e
+
     file = pq.ParquetFile(path)
     geo_meta = file.metadata.metadata.get(b"geo")
     if not geo_meta:
diff --git a/lonboard/_geoarrow/geopandas_interop.py b/lonboard/_geoarrow/geopandas_interop.py
@@ -16,7 +16,13 @@ def geopandas_to_geoarrow(
     columns: Optional[List[str]] = None,
     preserve_index: Optional[bool] = None,
 ) -> Table:
-    import pyarrow
+    try:
+        import pyarrow
+    except ImportError as e:
+        raise ImportError(
+            "pyarrow required for converting GeoPandas to arrow.\n"
+            "Run `pip install pyarrow`."
+        ) from e
 
     df_attr = gdf.drop(columns=[gdf._geometry_column_name])
 
diff --git a/lonboard/_layer.py b/lonboard/_layer.py
@@ -24,6 +24,7 @@
 import ipywidgets
 import traitlets
 from arro3.core import Table
+from arro3.core.types import ArrowStreamExportable
 
 from lonboard._base import BaseExtension, BaseWidget
 from lonboard._constants import EXTENSION_NAME, OGC_84
@@ -38,12 +39,11 @@
 from lonboard._utils import auto_downcast as _auto_downcast
 from lonboard._utils import get_geometry_column_index, remove_extension_kwargs
 from lonboard.traits import (
+    ArrowTableTrait,
     ColorAccessor,
     FloatAccessor,
     NormalAccessor,
-    PyarrowTableTrait,
 )
-from lonboard.types.arrow import ArrowStreamExportable
 
 if TYPE_CHECKING:
     import geopandas as gpd
@@ -283,7 +283,7 @@ class BaseArrowLayer(BaseLayer):
 
     # The following traitlets **are** serialized to JS
 
-    table: traitlets.TraitType
+    table: ArrowTableTrait
 
     def __init__(
         self,
@@ -714,7 +714,7 @@ def from_duckdb(
 
     _layer_type = traitlets.Unicode("column").tag(sync=True)
 
-    table = PyarrowTableTrait(allowed_geometry_types={EXTENSION_NAME.POINT})
+    table = ArrowTableTrait(allowed_geometry_types={EXTENSION_NAME.POINT})
     """A GeoArrow table with a Point or MultiPoint column.
 
     This is the fastest way to plot data from an existing GeoArrow source, such as
@@ -1001,7 +1001,7 @@ def from_duckdb(
 
     _layer_type = traitlets.Unicode("polygon").tag(sync=True)
 
-    table = PyarrowTableTrait(
+    table = ArrowTableTrait(
         allowed_geometry_types={EXTENSION_NAME.POLYGON, EXTENSION_NAME.MULTIPOLYGON}
     )
     """A GeoArrow table with a Polygon or MultiPolygon column.
@@ -1246,7 +1246,7 @@ def from_duckdb(
 
     _layer_type = traitlets.Unicode("scatterplot").tag(sync=True)
 
-    table = PyarrowTableTrait(
+    table = ArrowTableTrait(
         allowed_geometry_types={EXTENSION_NAME.POINT, EXTENSION_NAME.MULTIPOINT}
     )
     """A GeoArrow table with a Point or MultiPoint column.
@@ -1488,7 +1488,7 @@ def from_duckdb(
 
     _layer_type = traitlets.Unicode("path").tag(sync=True)
 
-    table = PyarrowTableTrait(
+    table = ArrowTableTrait(
         allowed_geometry_types={
             EXTENSION_NAME.LINESTRING,
             EXTENSION_NAME.MULTILINESTRING,
@@ -1659,7 +1659,7 @@ def from_duckdb(
 
     _layer_type = traitlets.Unicode("point-cloud").tag(sync=True)
 
-    table = PyarrowTableTrait(
+    table = ArrowTableTrait(
         allowed_geometry_types={EXTENSION_NAME.POINT}, allowed_dimensions={3}
     )
     """A GeoArrow table with a Point column.
@@ -1792,7 +1792,7 @@ def from_duckdb(
 
     _layer_type = traitlets.Unicode("solid-polygon").tag(sync=True)
 
-    table = PyarrowTableTrait(
+    table = ArrowTableTrait(
         allowed_geometry_types={EXTENSION_NAME.POLYGON, EXTENSION_NAME.MULTIPOLYGON}
     )
     """A GeoArrow table with a Polygon or MultiPolygon column.
@@ -1971,7 +1971,7 @@ def from_duckdb(
 
     _layer_type = traitlets.Unicode("heatmap").tag(sync=True)
 
-    table = PyarrowTableTrait(allowed_geometry_types={EXTENSION_NAME.POINT})
+    table = ArrowTableTrait(allowed_geometry_types={EXTENSION_NAME.POINT})
     """A GeoArrow table with a Point column.
 
     This is the fastest way to plot data from an existing GeoArrow source, such as
diff --git a/lonboard/_testing.py b/lonboard/_testing.py
diff --git a/lonboard/_viz.py b/lonboard/_viz.py
@@ -37,9 +37,9 @@
     import pyarrow
     import shapely.geometry
     import shapely.geometry.base
+    from arro3.core.types import ArrowArrayExportable, ArrowStreamExportable
     from numpy.typing import NDArray
 
-    from lonboard.types.arrow import ArrowArrayExportable, ArrowStreamExportable
     from lonboard.types.layer import (
         PathLayerKwargs,
         PolygonLayerKwargs,
@@ -372,8 +372,21 @@ def _viz_shapely_array(
 def _viz_geo_interface(
     data: dict, **kwargs
 ) -> List[Union[ScatterplotLayer, PathLayer, PolygonLayer]]:
-    import pyarrow as pa
-    import shapely
+    try:
+        import pyarrow as pa
+    except ImportError as e:
+        raise ImportError(
+            "pyarrow required for visualizing __geo_interface__ objects.\n"
+            "Run `pip install pyarrow`."
+        ) from e
+
+    try:
+        import shapely
+    except ImportError as e:
+        raise ImportError(
+            "shapely required for visualizing __geo_interface__ objects.\n"
+            "Run `pip install shapely`."
+        ) from e
 
     if data["type"] in [
         "Point",
diff --git a/lonboard/colormap.py b/lonboard/colormap.py
@@ -153,8 +153,14 @@ def apply_categorical_cmap(
             dimension will have a length of either `3` if `alpha` is `None`, or `4` is
             each color has an alpha value.
     """
-    import pyarrow as pa
-    import pyarrow.compute as pc
+    try:
+        import pyarrow as pa
+        import pyarrow.compute as pc
+    except ImportError as e:
+        raise ImportError(
+            "pyarrow required for apply_categorical_cmap.\n"
+            "Run `pip install pyarrow`."
+        ) from e
 
     # Import from PyCapsule interface
     if hasattr(values, "__arrow_c_array__"):
diff --git a/lonboard/experimental/_layer.py b/lonboard/experimental/_layer.py
@@ -10,10 +10,10 @@
 from lonboard._constants import EXTENSION_NAME
 from lonboard._layer import BaseArrowLayer
 from lonboard.traits import (
+    ArrowTableTrait,
     ColorAccessor,
     FloatAccessor,
     PointAccessor,
-    PyarrowTableTrait,
     TextAccessor,
 )
 
@@ -23,7 +23,7 @@ class ArcLayer(BaseArrowLayer):
 
     _layer_type = traitlets.Unicode("arc").tag(sync=True)
 
-    table = PyarrowTableTrait()
+    table = ArrowTableTrait()
     """A GeoArrow table.
 
     This is the fastest way to plot data from an existing GeoArrow source, such as
@@ -128,7 +128,7 @@ class TextLayer(BaseArrowLayer):
 
     _layer_type = traitlets.Unicode("text").tag(sync=True)
 
-    table = PyarrowTableTrait(allowed_geometry_types={EXTENSION_NAME.POINT})
+    table = ArrowTableTrait(allowed_geometry_types={EXTENSION_NAME.POINT})
     """A GeoArrow table with a Point or MultiPoint column.
 
     This is the fastest way to plot data from an existing GeoArrow source, such as
diff --git a/lonboard/traits.py b/lonboard/traits.py
@@ -137,7 +137,7 @@ def error(
                 raise TraitError(e)
 
 
-class PyarrowTableTrait(FixedErrorTraitType):
+class ArrowTableTrait(FixedErrorTraitType):
     """A trait to validate input for a geospatial Arrow-backed table
 
     Allowed input includes:
@@ -152,7 +152,10 @@ class from
     """
 
     default_value = None
-    info_text = "a pyarrow or GeoArrow Table"
+    info_text = (
+        "a table-like Arrow object, such as a pyarrow or arro3 Table or "
+        "RecordBatchReader"
+    )
 
     def __init__(
         self: TraitType,
diff --git a/lonboard/types/arrow.py b/lonboard/types/arrow.py
diff --git a/lonboard/types/layer.py b/lonboard/types/layer.py
@@ -4,7 +4,6 @@
 from typing import (
     List,
     Literal,
-    Protocol,
     Sequence,
     Tuple,
     Union,
@@ -13,6 +12,7 @@
 import numpy as np
 import pandas as pd
 import pyarrow
+from arro3.core.types import ArrowArrayExportable, ArrowStreamExportable
 from numpy.typing import NDArray
 
 from lonboard._base import BaseExtension
@@ -23,12 +23,6 @@
     from typing_extensions import TypedDict
 
 
-class ArrowArrayExportable(Protocol):
-    def __arrow_c_array__(
-        self, requested_schema: object | None = None
-    ) -> Tuple[object, object]: ...
-
-
 IntFloat = Union[int, float]
 Units = Literal["meters", "common", "pixels"]
 
@@ -41,6 +35,7 @@ def __arrow_c_array__(
     pyarrow.FixedSizeListArray,
     pyarrow.ChunkedArray,
     ArrowArrayExportable,
+    ArrowStreamExportable,
 ]
 FloatAccessorInput = Union[
     int,
@@ -50,6 +45,7 @@ def __arrow_c_array__(
     pyarrow.FloatingPointArray,
     pyarrow.ChunkedArray,
     ArrowArrayExportable,
+    ArrowStreamExportable,
 ]
 NormalAccessorInput = Union[
     List[int],
@@ -59,6 +55,7 @@ def __arrow_c_array__(
     pyarrow.FixedSizeListArray,
     pyarrow.ChunkedArray,
     ArrowArrayExportable,
+    ArrowStreamExportable,
 ]
 TextAccessorInput = Union[
     str,
@@ -68,6 +65,7 @@ def __arrow_c_array__(
     pyarrow.LargeStringArray,
     pyarrow.ChunkedArray,
     ArrowArrayExportable,
+    ArrowStreamExportable,
 ]
 
 
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -180,6 +180,7 @@ plugins:
             - https://geopy.readthedocs.io/en/stable/objects.inv
             - https://ipython.readthedocs.io/en/stable/objects.inv
             - https://ipywidgets.readthedocs.io/en/stable/objects.inv
+            - https://kylebarron.dev/arro3/latest/objects.inv
             - https://matplotlib.org/stable/objects.inv
             - https://numpy.org/doc/stable/objects.inv
             - https://pandas.pydata.org/pandas-docs/stable/objects.inv
diff --git a/tests/test_traits.py b/tests/test_traits.py
@@ -11,10 +11,10 @@
 from lonboard._layer import BaseArrowLayer, BaseLayer
 from lonboard.layer_extension import DataFilterExtension
 from lonboard.traits import (
+    ArrowTableTrait,
     ColorAccessor,
     FloatAccessor,
     NormalAccessor,
-    PyarrowTableTrait,
 )
 
 
@@ -176,7 +176,7 @@ class FilterValueAccessorWidget(BaseArrowLayer):
         sync=True, **ipywidgets.widget_serialization
     )
 
-    table = PyarrowTableTrait()
+    table = ArrowTableTrait()
 
     def __init__(self, *args, **kwargs):
         # Any tests that are intended to pass validation checks must also have 3 rows,
