feat: Support passing ChunkedArray/ArrayReader into viz (#824)

### Change list

Currently, passing a `ChunkedArray`/`ArrayReader` into `viz` fails
because any `ArrowStreamExportable` it tries to import as a table.

This adds a helper for checking the field of an input stream. If the
field is struct but **not** a GeoArrow extension type, then it assumes
the input is a `Table`. We need to check for the extension type because
that ensures that we can render a struct-encoded point column or a
`geoarrow.box` column.

Also:

- Update `test_viz.py` for latest geoarrow-rust-core version 0.4.
- Rename `geoarrow` test folder so that when starting a Python kernel in
the above directory it doesn't search in the test files for the
`geoarrow` python library.
- Update some other dev deps

#823 will be a follow
up

Author: kylebarron

lonboard/_geoarrow/c_stream_import.py

@@ -0,0 +1,32 @@
+"""Import Arrow data from a C stream into a Table."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from arro3.core import ArrayReader, ChunkedArray, DataType, Table
+
+if TYPE_CHECKING:
+    from arro3.core.types import ArrowStreamExportable
+
+
+def import_arrow_c_stream(stream: ArrowStreamExportable) -> Table | ChunkedArray:
+    """Import an Arrow stream into a Table.
+
+    If the stream is a struct but **not** a geoarrow extension type, then we assume it's
+    table input. Otherwise, we assume it's an array input, which we then coerce into a
+    Table with a single field named "geometry".
+    """
+    array_reader = ArrayReader.from_arrow(stream)
+
+    if DataType.is_struct(
+        array_reader.field.type,
+    ) and not array_reader.field.metadata_str.get(
+        "ARROW:extension:name",
+        "",
+    ).startswith("geoarrow"):
+        # If the field is a struct but **not** a geoarrow extension type, then we assume
+        # it's table input.
+        return Table.from_arrow(array_reader)
+
+    return ChunkedArray.from_arrow(array_reader)

lonboard/_viz.py

@@ -20,6 +20,7 @@
 
 from lonboard._compat import check_pandas_version
 from lonboard._constants import EXTENSION_NAME
+from lonboard._geoarrow.c_stream_import import import_arrow_c_stream
 from lonboard._geoarrow.extension_types import construct_geometry_array
 from lonboard._geoarrow.geopandas_interop import geopandas_to_geoarrow
 from lonboard._geoarrow.parse_wkb import parse_serialized_table
@@ -275,12 +276,19 @@ def create_layers_from_data_input(
     # Anything with __arrow_c_array__
     if hasattr(data, "__arrow_c_array__"):
         data = cast("ArrowArrayExportable", data)
-        return _viz_geoarrow_array(data, **kwargs)
+        array = Array.from_arrow(data)
+        ca = ChunkedArray([array])
+        return _viz_geoarrow_chunked_array(ca, **kwargs)
 
     # Anything with __arrow_c_stream__
     if hasattr(data, "__arrow_c_stream__"):
         data = cast("ArrowStreamExportable", data)
-        return _viz_geoarrow_table(Table.from_arrow(data), **kwargs)
+        imported_stream = import_arrow_c_stream(data)
+        if isinstance(imported_stream, Table):
+            return _viz_geoarrow_table(imported_stream, **kwargs)
+
+        assert isinstance(imported_stream, ChunkedArray)
+        return _viz_geoarrow_chunked_array(imported_stream, **kwargs)
 
     # Anything with __geo_interface__
     if hasattr(data, "__geo_interface__"):
@@ -432,16 +440,15 @@ def _viz_geo_interface(
     raise ValueError(f"type '{geo_interface_type}' not supported.")
 
 
-def _viz_geoarrow_array(
-    data: ArrowArrayExportable,
+def _viz_geoarrow_chunked_array(
+    ca: ChunkedArray,
     **kwargs: Any,
 ) -> list[ScatterplotLayer | PathLayer | PolygonLayer]:
-    array = Array.from_arrow(data)
-    field = array.field.with_name("geometry")
+    field = ca.field.with_name("geometry")
     schema = Schema([field])
-    table = Table.from_arrays([array], schema=schema)
+    table = Table.from_arrays([ca], schema=schema)
 
-    num_rows = len(array)
+    num_rows = len(ca)
     if num_rows <= np.iinfo(np.uint8).max:
         arange_col = Array(np.arange(num_rows, dtype=np.uint8))
     elif num_rows <= np.iinfo(np.uint16).max:
@@ -575,4 +582,4 @@ def _viz_geoarrow_table(
 
         return [PolygonLayer(table=table, **polygon_kwargs)]
 
-    raise ValueError(f"Unsupported extension type: '{geometry_ext_type}'.")
+    raise ValueError(f"Unsupported extension type: '{geometry_ext_type!r}'.")

pyproject.toml

@@ -72,20 +72,21 @@ docs = [
     "black>=24.8.0",
 ]
 dev = [
-    "duckdb>=1.2.0",
+    "duckdb>=1.3.0",
     "fiona<1.10",
     "geoarrow-pyarrow>=0.2.0",
-    "geoarrow-rust-core>=0.3.0",
+    "geoarrow-rust-core>=0.4.0",
     "geodatasets>=2024.8.0",
     "jupyterlab>=4.3.3",
     "matplotlib>=3.7.5",
     "movingpandas>=0.20.0",
     "palettable>=3.3.3",
+    "pandas-stubs>=2.2.2.240807",
     "pre-commit>=3.5.0",
     "pyarrow>=17.0.0",
     "pyogrio>=0.9.0",
     "pytest>=8.3.4",
-    "ruff>=0.11.0",
+    "ruff>=0.12.0",
     "sidecar>=0.7.0",
     "types-shapely>=2.1.0.20250512",
 ]

tests/geoarrow/__init__.py renamed to tests/geoarrow_tests/__init__.py

@@ -1,16 +1,15 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import TYPE_CHECKING, cast
 
 import geoarrow.pyarrow as gap
 import geodatasets
 import numpy as np
 import pyarrow as pa
 import pyarrow.parquet as pq
 import pytest
-from geoarrow.rust.core import geometry_col, read_pyogrio, to_wkb
-from pyogrio.raw import read_arrow
+from geoarrow.rust.core import to_wkb
+from pyogrio import read_arrow
 
 from lonboard import PathLayer, PolygonLayer, ScatterplotLayer, viz
 from lonboard._constants import EXTENSION_NAME
@@ -19,9 +18,6 @@
 
 fixtures_dir = Path(__file__).parent / "fixtures"
 
-if TYPE_CHECKING:
-    from arro3.core import Table
-
 
 def mixed_shapely_geoms():
     shapely = pytest.importorskip("shapely")
@@ -156,31 +152,31 @@ def test_viz_shapely_mixed_array():
     assert isinstance(map_.layers[2], PolygonLayer)
 
 
-# read_pyogrio currently keeps geometries as WKB
 @pytest.mark.skipif(not compat.HAS_SHAPELY, reason="shapely not available")
-def test_viz_geoarrow_rust_table():
-    table = read_pyogrio(geodatasets.get_path("naturalearth.land"))
+def test_viz_pyogrio_table():
+    _meta, table = read_arrow(geodatasets.get_path("naturalearth.land"))
     map_ = viz(table)
     assert isinstance(map_.layers[0], PolygonLayer)
 
 
-# read_pyogrio currently keeps geometries as WKB
 @pytest.mark.skipif(not compat.HAS_SHAPELY, reason="shapely not available")
-def test_viz_geoarrow_rust_array():
-    # `read_pyogrio` has incorrect typing currently
-    # https://github.com/geoarrow/geoarrow-rs/pull/807
-    table = cast("Table", read_pyogrio(geodatasets.get_path("naturalearth.land")))
-    map_ = viz(geometry_col(table).chunk(0))
+def test_viz_pyogrio_chunked_array():
+    _meta, table = read_arrow(geodatasets.get_path("naturalearth.land"))
+    map_ = viz(table["wkb_geometry"])
+    assert isinstance(map_.layers[0], PolygonLayer)
+
+
+@pytest.mark.skipif(not compat.HAS_SHAPELY, reason="shapely not available")
+def test_viz_pyogrio_array():
+    _meta, table = read_arrow(geodatasets.get_path("naturalearth.land"))
+    map_ = viz(table["wkb_geometry"].chunk(0))
     assert isinstance(map_.layers[0], PolygonLayer)
 
 
-# read_pyogrio currently keeps geometries as WKB
-@pytest.mark.skip("to_wkb gives panic on todo!(), not yet implemented")
 @pytest.mark.skipif(not compat.HAS_SHAPELY, reason="shapely not available")
-def test_viz_geoarrow_rust_wkb_array():
-    table = read_pyogrio(geodatasets.get_path("naturalearth.land"))
-    arr = geometry_col(table).chunk(0)
-    wkb_arr = to_wkb(arr)
+def test_viz_array_reader():
+    _meta, table = read_arrow(geodatasets.get_path("naturalearth.land"))
+    wkb_arr = to_wkb(table["wkb_geometry"])
     map_ = viz(wkb_arr)
     assert isinstance(map_.layers[0], PolygonLayer)