fix: map geoarrow to geoarrow-types

marimo/_plugins/ui/_impl/tables/selection.py

@@ -19,6 +19,14 @@ def add_selection_column(data: T) -> tuple[T, bool]:
                 native = df.to_native().copy()
                 native.insert(0, INDEX_COLUMN_NAME, range(len(native)))
                 return cast(T, native), True
+            if df.implementation.is_pyarrow():
+                import pyarrow as pa
+
+                native = df.to_native()
+                index = pa.array(range(len(native)))
+                return cast(
+                    T, native.add_column(0, INDEX_COLUMN_NAME, index)
+                ), True
             return df.with_row_index(name=INDEX_COLUMN_NAME).to_native(), True  # type: ignore[return-value]
         return data, True  # already has a row index
     return data, False

marimo/_runtime/packages/module_name_to_pypi_name.py

@@ -446,6 +446,7 @@ def module_name_to_pypi_name() -> dict[str, str]:
         "gen_3dwallet": "3d-wallet-generator",
         "gendimen": "android-gendimen",
         "genshi": "Genshi",
+        "geoarrow": "geoarrow-types",
         "geonode": "GeoNode",
         "geoserver": "gsconfig",
         "geraldo": "Geraldo",

tests/_plugins/ui/_impl/tables/snapshots/false_positives.pandas.field_types.json

@@ -10,7 +10,7 @@
     "wkt_looking",
     [
       "string",
-      "str"
+      "object"
     ]
   ],
   [

tests/_plugins/ui/_impl/tables/test_geometry.py

@@ -3,6 +3,10 @@
 from __future__ import annotations
 
 import json
+import subprocess
+import sys
+from contextlib import contextmanager
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
 
 import pytest
@@ -272,6 +276,24 @@ def test_formatting_preserves_schema_metadata(self) -> None:
             "geoarrow.wkb",
         )
 
+    def test_selection_column_preserves_geoarrow_metadata(self) -> None:
+        from marimo._plugins.ui._impl.tables.selection import (
+            INDEX_COLUMN_NAME,
+            add_selection_column,
+        )
+
+        table, has_stable_row_id = add_selection_column(
+            geo.arrow_wkb_known_crs()
+        )
+        field = table.schema.field("geom")
+
+        assert has_stable_row_id is True
+        assert INDEX_COLUMN_NAME in table.column_names
+        assert field.metadata[b"ARROW:extension:name"] == b"geoarrow.wkb"
+        assert ui.table(geo.arrow_wkb_known_crs())._manager.get_field_type(
+            "geom"
+        ) == ("geometry", "geoarrow.wkb")
+
     def test_formatting_skips_geometry_columns_in_mapping(self) -> None:
         manager = get_table_manager(geo.arrow_wkb_known_crs())
 
@@ -684,6 +706,170 @@ def test_filter_is_ignored(self) -> None:
         )
 
 
+@contextmanager
+def _closing_data(make_data: Any) -> Any:
+    produced = make_data()
+    if isinstance(produced, tuple):
+        conn, data = produced
+        try:
+            yield data
+        finally:
+            conn.close()
+        return
+    yield produced
+
+
+def _duckdb_geometry() -> tuple[Any, Any]:
+    conn = geo.duckdb_spatial_connection()
+    return conn, geo.duckdb_geometry_relation(conn)
+
+
+def _duckdb_crs_geometry() -> tuple[Any, Any]:
+    conn = geo.duckdb_crs_geometry_connection()
+    return conn, conn.table("crs_geometry")
+
+
+@pytest.mark.parametrize(
+    ("make_data", "expected_geometry"),
+    [
+        pytest.param(
+            geo.gdf_point_known_crs,
+            {"geometry": ("geometry", "geometry")},
+            marks=pytest.mark.requires("geopandas"),
+            id="gdf_point_known_crs",
+        ),
+        pytest.param(
+            geo.gdf_point_missing_crs,
+            {"geometry": ("geometry", "geometry")},
+            marks=pytest.mark.requires("geopandas"),
+            id="gdf_point_missing_crs",
+        ),
+        pytest.param(
+            geo.gdf_multi_geometry,
+            {
+                "geom_a": ("geometry", "geometry"),
+                "geom_b": ("geometry", "geometry"),
+            },
+            marks=pytest.mark.requires("geopandas"),
+            id="gdf_multi_geometry",
+        ),
+        pytest.param(
+            geo.gdf_no_active_geometry,
+            {"g": ("geometry", "geometry")},
+            marks=pytest.mark.requires("geopandas"),
+            id="gdf_no_active_geometry",
+        ),
+        pytest.param(
+            geo.gdf_stale_pointer,
+            {"geom": ("geometry", "geometry")},
+            marks=pytest.mark.requires("geopandas"),
+            id="gdf_stale_pointer",
+        ),
+        pytest.param(
+            geo.gdf_dropped_active,
+            {"geom_b": ("geometry", "geometry")},
+            marks=pytest.mark.requires("geopandas"),
+            id="gdf_dropped_active",
+        ),
+        pytest.param(
+            geo.gdf_with_null,
+            {"geometry": ("geometry", "geometry")},
+            marks=pytest.mark.requires("geopandas"),
+            id="gdf_with_null",
+        ),
+        pytest.param(
+            geo.gdf_all_null,
+            {"geometry": ("geometry", "geometry")},
+            marks=pytest.mark.requires("geopandas"),
+            id="gdf_all_null",
+        ),
+        pytest.param(
+            geo.gdf_mixed_types,
+            {"geometry": ("geometry", "geometry")},
+            marks=pytest.mark.requires("geopandas"),
+            id="gdf_mixed_types",
+        ),
+        pytest.param(
+            geo.gdf_3d,
+            {"geometry": ("geometry", "geometry")},
+            marks=pytest.mark.requires("geopandas"),
+            id="gdf_3d",
+        ),
+        pytest.param(
+            geo.arrow_wkb_known_crs,
+            {"geom": ("geometry", "geoarrow.wkb")},
+            marks=pytest.mark.requires("pyarrow"),
+            id="arrow_wkb_known_crs",
+        ),
+        pytest.param(
+            geo.arrow_wkt,
+            {"geom": ("geometry", "geoarrow.wkt")},
+            marks=pytest.mark.requires("pyarrow"),
+            id="arrow_wkt",
+        ),
+        pytest.param(
+            geo.arrow_wkb_missing_crs,
+            {"geom": ("geometry", "geoarrow.wkb")},
+            marks=pytest.mark.requires("pyarrow"),
+            id="arrow_wkb_missing_crs",
+        ),
+        pytest.param(
+            geo.arrow_ogc_wkb,
+            {"geom": ("geometry", "ogc.wkb")},
+            marks=pytest.mark.requires("pyarrow"),
+            id="arrow_ogc_wkb",
+        ),
+        pytest.param(
+            geo.arrow_other_geoarrow,
+            {"geom": ("geometry", "geoarrow.point")},
+            marks=pytest.mark.requires("pyarrow"),
+            id="arrow_other_geoarrow",
+        ),
+        pytest.param(
+            _duckdb_geometry,
+            {"geom": ("geometry", "GEOMETRY")},
+            marks=pytest.mark.requires("duckdb", "pyarrow"),
+            id="duckdb_geometry",
+        ),
+        pytest.param(
+            _duckdb_crs_geometry,
+            {"geom": ("geometry", "GEOMETRY('OGC:CRS84')")},
+            marks=pytest.mark.requires("duckdb", "pyarrow"),
+            id="duckdb_crs_geometry",
+        ),
+    ],
+)
+class TestHostCorpusSmoke:
+    """Host-level smoke over the geometry fixture corpus.
+
+    Each fixture is wrapped in ui.table. Search and column summaries must
+    return clean payloads with the expected geometry field types.
+    """
+
+    def test_search_and_summaries(
+        self,
+        make_data: Any,
+        expected_geometry: dict[str, tuple[str, str]],
+    ) -> None:
+        with _closing_data(make_data) as data:
+            table = ui.table(data, show_column_summaries=True)
+            field_types = dict(table._manager.get_field_types())
+            for name, expected in expected_geometry.items():
+                assert field_types[name] == expected
+
+            response = table._search(
+                SearchTableArgs(page_size=10, page_number=0)
+            )
+            rows = json.loads(response.data)
+            assert isinstance(rows, list)
+
+            summaries = table._get_column_summaries(ColumnSummariesArgs())
+            for name in expected_geometry:
+                stats = summaries.stats[name]
+                assert stats.unique is None
+                assert stats.min is None
+
+
 class TestManagerTemplateHook:
     @pytest.mark.requires("polars")
     def test_polars_uses_semantic_type(self) -> None:
@@ -714,3 +900,64 @@ def test_ibis_uses_semantic_type(self) -> None:
         }
 
         assert manager.get_field_type("geometry") == ("geometry", "geometry")
+
+
+@pytest.mark.requires("pyarrow")
+def test_geoarrow_wkb_without_geopandas_or_shapely() -> None:
+    repo_root = str(Path(__file__).resolve().parents[5])
+    script = """
+import json
+import sys
+
+sys.path.insert(0, sys.argv[1])
+
+
+class BlockedDependency:
+    def find_spec(self, fullname, path=None, target=None):
+        del path, target
+        if fullname.partition(".")[0] in {"geopandas", "shapely"}:
+            raise RuntimeError(f"blocked import: {fullname}")
+        return None
+
+
+sys.meta_path.insert(0, BlockedDependency())
+
+import pyarrow as pa
+from marimo._plugins.ui._impl.tables.utils import get_table_manager
+
+wkb = bytes.fromhex("0101000000000000000000f03f0000000000000040")
+schema = pa.schema(
+    [
+        pa.field("a", pa.int64()),
+        pa.field(
+            "geom",
+            pa.binary(),
+            metadata={
+                b"ARROW:extension:name": b"geoarrow.wkb",
+                b"ARROW:extension:metadata": b'{"crs": "EPSG:4326"}',
+            },
+        ),
+    ]
+)
+table = pa.table({"a": [0, 1], "geom": [wkb, None]}, schema=schema)
+manager = get_table_manager(table)
+
+assert manager.get_field_type("geom") == ("geometry", "geoarrow.wkb")
+rows = json.loads(manager.to_json_str())
+assert rows[0]["geom"] == f"<geometry, {len(wkb)} B>"
+assert rows[1]["geom"] is None
+stats = manager.get_stats("geom")
+assert stats.total == 2
+assert stats.nulls == 1
+assert stats.unique is None
+assert stats.min is None
+"""
+    try:
+        subprocess.run(
+            [sys.executable, "-c", script, repo_root],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except subprocess.CalledProcessError as e:
+        raise AssertionError(e.stderr) from e

tests/_plugins/ui/_impl/tables/test_geometry_fixtures.py

@@ -3,21 +3,38 @@
 from __future__ import annotations
 
 import json
-import os
+import re
 
 import pytest
 
 from marimo._plugins.ui._impl.tables.utils import get_table_manager
 from tests._plugins.ui._impl.tables import geometry_fixtures as geo
 from tests.mocks import snapshotter
 
-pytestmark = pytest.mark.skipif(
-    os.getenv("CI") is not None,
-    reason="Geometry characterization tests are disabled until they can run reliably in CI.",
-)
-
 snapshot = snapshotter(__file__)
 
+_PANDAS_STRING_STORAGE = frozenset({"str", "object"})
+
+
+def _assert_field_types(
+    actual: list[tuple[str, tuple[str, str]]],
+    expected: list[tuple[str, tuple[str, str]]],
+) -> None:
+    """Compare field types, allowing pandas 2/3 string storage names."""
+    assert len(actual) == len(expected)
+    for (name, (sem, ext)), (exp_name, (exp_sem, exp_ext)) in zip(
+        actual, expected, strict=True
+    ):
+        assert name == exp_name
+        assert sem == exp_sem
+        if sem == "string" and {ext, exp_ext} <= _PANDAS_STRING_STORAGE:
+            continue
+        assert ext == exp_ext
+
+
+def _normalize_pandas_string_storage(text: str) -> str:
+    return re.sub(r'\["string",\s*"str"\]', '["string", "object"]', text)
+
 
 @pytest.mark.requires("geopandas")
 class TestGeoPandasCharacterization:
@@ -30,10 +47,13 @@ def test_preserves_geodataframe_on_ingest(self) -> None:
 
     def test_geometry_column_types_geometry(self) -> None:
         manager = get_table_manager(geo.gdf_point_known_crs())
-        assert dict(manager.get_field_types()) == {
-            "name": ("string", "str"),
-            "geometry": ("geometry", "geometry"),
-        }
+        _assert_field_types(
+            manager.get_field_types(),
+            [
+                ("name", ("string", "str")),
+                ("geometry", ("geometry", "geometry")),
+            ],
+        )
 
     @pytest.mark.parametrize(
         ("make_frame", "expected_field_types"),
@@ -104,7 +124,7 @@ def test_corpus_loads_and_serializes(
         self, make_frame, expected_field_types
     ) -> None:
         manager = get_table_manager(make_frame())
-        assert manager.get_field_types() == expected_field_types
+        _assert_field_types(manager.get_field_types(), expected_field_types)
         assert isinstance(manager.to_json_str(), str)
 
 
@@ -196,7 +216,9 @@ def test_pandas(self) -> None:
         manager = get_table_manager(pd.DataFrame(geo.false_positive_data()))
         snapshot(
             "false_positives.pandas.field_types.json",
-            json.dumps(manager.get_field_types()),
+            _normalize_pandas_string_storage(
+                json.dumps(manager.get_field_types())
+            ),
         )
         snapshot(
             "false_positives.pandas.json",