diff --git a/marimo/_plugins/ui/_impl/tables/selection.py b/marimo/_plugins/ui/_impl/tables/selection.py index a91dd8bc69b..f79bfa08f91 100644 --- a/marimo/_plugins/ui/_impl/tables/selection.py +++ b/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 diff --git a/marimo/_runtime/packages/module_name_to_pypi_name.py b/marimo/_runtime/packages/module_name_to_pypi_name.py index d4ef2b70856..10b43b1a87f 100644 --- a/marimo/_runtime/packages/module_name_to_pypi_name.py +++ b/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", diff --git a/tests/_plugins/ui/_impl/tables/snapshots/false_positives.pandas.field_types.json b/tests/_plugins/ui/_impl/tables/snapshots/false_positives.pandas.field_types.json index 3504b631ef8..f13d8b4e5c0 100644 --- a/tests/_plugins/ui/_impl/tables/snapshots/false_positives.pandas.field_types.json +++ b/tests/_plugins/ui/_impl/tables/snapshots/false_positives.pandas.field_types.json @@ -10,7 +10,7 @@ "wkt_looking", [ "string", - "str" + "object" ] ], [ diff --git a/tests/_plugins/ui/_impl/tables/test_geometry.py b/tests/_plugins/ui/_impl/tables/test_geometry.py index f95a475b88a..35bf26542b4 100644 --- a/tests/_plugins/ui/_impl/tables/test_geometry.py +++ b/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"" +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 diff --git a/tests/_plugins/ui/_impl/tables/test_geometry_fixtures.py b/tests/_plugins/ui/_impl/tables/test_geometry_fixtures.py index 8992a47493a..c6f142fd311 100644 --- a/tests/_plugins/ui/_impl/tables/test_geometry_fixtures.py +++ b/tests/_plugins/ui/_impl/tables/test_geometry_fixtures.py @@ -3,7 +3,7 @@ from __future__ import annotations import json -import os +import re import pytest @@ -11,13 +11,30 @@ 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",