Support geometry columns end-to-end so GeoDataFrames render as maps (#1903)

Author: ghostiee-11

lumen/ai/agents/deck_gl.py

@@ -96,6 +96,7 @@ class DeckGLAgent(BaseCodeAgent):
             "Use for 3D geographic visualizations, map-based data, or when user requests DeckGL/deck.gl",
             "Use for large-scale geospatial data with latitude/longitude coordinates",
             "Use for hexbin aggregations, heatmaps, or 3D extruded visualizations on maps",
+            "Can render a geometry column (GeoDataFrame polygons/lines/points) on a basemap, suited to 3D/extruded maps or geometry whose schema is 'geographic' (a lat/lon CRS) so the basemap gives real-world context",
         ]
     )
 

lumen/ai/agents/hvplot.py

@@ -19,6 +19,10 @@ class hvPlotAgent(BaseViewAgent):
         default=[
             "Use for exploratory data analysis, interactive plots, and dynamic filtering",
             "Use for quick, iterative data visualization during analysis",
+            "Can render a geometry column (GeoDataFrame polygons/lines/points) as a "
+            "2D choropleth, shading shapes by a value column without a basemap; fits "
+            "geometry whose schema is not 'geographic' (a projected or absent CRS) or "
+            "when a basemap adds no context",
         ]
     )
 

lumen/ai/controls/ingest/utils.py

@@ -326,7 +326,12 @@ def read_geo_file(
 
     return FileReadResult(
         tables={alias: df},
-        source_params={"initializers": ["INSTALL spatial;\nLOAD spatial;"]},
+        source_params={
+            "initializers": ["INSTALL spatial;\nLOAD spatial;"],
+            # DuckDB stores geometry without a CRS, so carry the source CRS
+            # through for DuckDBSource to reapply after the WKB roundtrip.
+            "geometry_crs": str(geo_df.crs) if geo_df.crs is not None else None,
+        },
         conversions={alias: conversion},
     )
 

lumen/ai/prompts/hvPlotAgent/main.jinja2

@@ -4,5 +4,6 @@
 Generate the plot the user requested. Note that `x`, `y`, `by` and `groupby` fields MUST ALL be unique columns;
 no repeated columns are allowed. Do not arbitrarily set `groupby` and `by` fields unless explicitly requested. If a
 histogram is requested, use `y` instead of `x`. If x is categorical or strings, prefer barh over bar, and use `y` for
-the values.
+the values. If a column has `format: geometry`, set `kind` from its `geometry_type`: `polygons` for Polygon or
+MultiPolygon, `paths` for LineString or MultiLineString, otherwise `points`.
 {% endblock %}

lumen/sources/duckdb.py

@@ -19,7 +19,9 @@
 from ..transforms.sql import (
     SQLCount, SQLFilter, SQLLimit, SQLSelectFrom,
 )
-from ..util import detect_file_encoding, normalize_table_name
+from ..util import (
+    detect_file_encoding, normalize_table_name, try_import_geopandas,
+)
 from .base import BaseSQLSource, Source, cached
 
 if TYPE_CHECKING:
@@ -66,6 +68,11 @@ class DuckDBSource(BaseSQLSource):
         Whether the data is ephemeral, i.e. manually inserted into the
         DuckDB table or derived from real data.""")
 
+    geometry_crs = param.String(default=None, allow_None=True, doc="""
+        CRS to reapply to geometry columns after the WKB roundtrip through
+        DuckDB, which stores geometry without a CRS. Populated from the source
+        data at ingest; may also be set explicitly for a known dataset.""")
+
     read_only = param.Boolean(default=None, doc="""
         Whether to open the DuckDB database in read-only mode.""")
 
@@ -500,11 +507,40 @@ def create_sql_expr_source(
         source._file_based_tables = {**self._file_based_tables, **source._file_based_tables}
         return source
 
+    def _fetch_df(
+        self, cursor, sql_expr: str, params: list | dict | None = None,
+        date_as_object: bool = False
+    ):
+        """Fetch a query result, safely handling native GEOMETRY columns.
+
+        DuckDB cannot convert a native GEOMETRY column to NumPy, so re-select
+        any geometry columns as WKB via ST_AsWKB before fetching, then rebuild
+        a GeoDataFrame when geopandas is available (WKB bytes otherwise).
+        """
+        rel = cursor.execute(sql_expr, params) if params else cursor.execute(sql_expr)
+        geom_cols = [d[0] for d in rel.description if str(d[1]) == 'GEOMETRY']
+        if not geom_cols:
+            return rel.fetch_df(date_as_object=date_as_object)
+
+        selected = ', '.join(
+            f'ST_AsWKB("{d[0]}"::GEOMETRY) AS "{d[0]}"'
+            if d[0] in geom_cols else f'"{d[0]}"'
+            for d in rel.description
+        )
+        wrapped = f'SELECT {selected} FROM ({sql_expr})'
+        rel = cursor.execute(wrapped, params) if params else cursor.execute(wrapped)
+        df = rel.fetch_df(date_as_object=date_as_object)
+        if gpd := try_import_geopandas():
+            for col in geom_cols:
+                df[col] = gpd.GeoSeries.from_wkb(
+                    df[col].apply(bytes), crs=self.geometry_crs
+                )
+            df = gpd.GeoDataFrame(df, geometry=geom_cols[0])
+        return df
+
     def execute(self, sql_query: str, params: list | dict | None = None, *args, **kwargs):
         with self._connection.cursor() as cursor:
-            if params:
-                return cursor.execute(sql_query, params, *args, **kwargs).fetch_df()
-            return cursor.execute(sql_query, *args, **kwargs).fetch_df()
+            return self._fetch_df(cursor, sql_query, params)
 
     def get_tables(self):
         if isinstance(self.tables, dict | list):
@@ -538,20 +574,9 @@ def get(self, table, **query):
 
         # Apply stored SQL parameters if available for this table
         with self._connection.cursor() as cursor:
-            if table in self.table_params:
-                rel = cursor.execute(sql_expr, self.table_params[table])
-            else:
-                rel = cursor.execute(sql_expr)
-            has_geom = any(d[0] == 'geometry' and d[1] == 'BINARY' for d in rel.description)
-            df = rel.fetch_df(date_as_object=True)
-            if has_geom:
-                import geopandas as gpd
-                geom_rel = cursor.execute(
-                    f'SELECT ST_AsWKB(geometry::GEOMETRY) as geometry FROM ({sql_expr})'
-                )
-                geom_df = geom_rel.fetch_df()
-                df['geometry'] = gpd.GeoSeries.from_wkb(geom_df.geometry.apply(bytes))
-                df = gpd.GeoDataFrame(df)
+            df = self._fetch_df(
+                cursor, sql_expr, self.table_params.get(table), date_as_object=True
+            )
         if not self.filter_in_sql:
             df = Filter.apply_to(df, conditions=conditions)
         return df

lumen/tests/ai/test_agents.py

@@ -20,6 +20,8 @@
     AnalysisAgent, ChatAgent, SQLAgent, VegaLiteAgent,
 )
 from lumen.ai.agents.analysis import make_analysis_model
+from lumen.ai.agents.deck_gl import DeckGLAgent
+from lumen.ai.agents.hvplot import hvPlotAgent
 from lumen.ai.agents.sql import make_sql_model
 from lumen.ai.agents.vega_lite import VegaLiteSpec, VegaLiteSpecUpdate
 from lumen.ai.analysis import Analysis
@@ -277,6 +279,13 @@ class ExtendedChatAgent(ChatAgent):
         assert "Footer appended." in prompt
 
 
+def test_map_agents_route_geometry_columns():
+    """hvPlot and DeckGL agents advertise a geometry-column condition so the
+    coordinator routes GeoDataFrame data to a map-capable view."""
+    assert any("geometry" in c.lower() for c in hvPlotAgent.conditions)
+    assert any("geometry" in c.lower() for c in DeckGLAgent.conditions)
+
+
 def test_sqlagent_active_filters_describes_conditions():
     """SQLAgent._active_filters turns the interactive slider filters into
     WHERE-style conditions (skipping inactive/full-range ones) so a follow-up

lumen/tests/ai/test_controls/test_upload.py

@@ -7,8 +7,16 @@
 import pytest
 
 from lumen.ai.controls import SourceResult, UploadedFileRow
+from lumen.ai.controls.ingest.utils import read_geo_file
 from lumen.sources.duckdb import DuckDBSource
 
+try:
+    import geopandas as gpd
+
+    from shapely.geometry import Polygon
+except ImportError:
+    gpd = None
+
 
 @pytest.mark.asyncio
 class TestDocumentVectorStoreIntegration:
@@ -363,3 +371,16 @@ class TestUploadControlsUX:
     def test_upload_button_label_is_explicit(self, upload_controls):
         """Upload controls should use explicit upload action text."""
         assert upload_controls._add_button.name == "Upload file(s)"
+
+
+def test_read_geo_file_captures_crs():
+    """read_geo_file surfaces the source CRS in source_params so DuckDBSource
+    can reapply it after the WKB roundtrip (gh-1904)."""
+    if gpd is None:
+        pytest.skip("geopandas is not installed")
+    gdf = gpd.GeoDataFrame(
+        {"name": ["a"]}, geometry=[Polygon([(0, 0), (1, 0), (1, 1)])], crs="EPSG:4326"
+    )
+    buf = io.BytesIO(gdf.to_json().encode())
+    result = read_geo_file(buf, "geojson", "geo")
+    assert result.source_params["geometry_crs"] == "EPSG:4326"

lumen/tests/sources/test_duckdb.py

@@ -18,6 +18,13 @@
 except ImportError:
     pytestmark = pytest.mark.skip(reason="Duckdb is not installed")
 
+try:
+    import geopandas as gpd
+
+    from shapely.geometry import Polygon
+except ImportError:
+    gpd = None
+
 
 @pytest.fixture
 def duckdb_file_source():
@@ -239,7 +246,10 @@ def test_duckdb_transforms_cache(duckdb_source, source_tables):
     assert cache_key in duckdb_source._cache
 
     expected = df_test_sql.groupby('B')['A'].sum().reset_index()
-    pd.testing.assert_frame_equal(duckdb_source._cache[cache_key], expected)
+    # DuckDB GROUP BY returns rows in a nondeterministic order; sort by the
+    # group key before comparing against pandas' sorted groupby output.
+    actual = duckdb_source._cache[cache_key].sort_values('B').reset_index(drop=True)
+    pd.testing.assert_frame_equal(actual, expected)
 
     cache_key = duckdb_source._get_key('test_sql', sql_transforms=transforms)
     assert cache_key in duckdb_source._cache
@@ -1345,3 +1355,92 @@ def test_read_only_create_sql_expr_source_expands_tables(duckdb_file_source):
 
     # Ensure new tables did NOT modify the original source
     assert source.execute("SHOW TABLES").shape[0] == 1  # Only 'test' table in original source
+
+
+def _spatial_source():
+    """Build an in-memory DuckDBSource holding a native GEOMETRY table.
+
+    Mirrors the ingest flow: WKB bytes -> ST_GeomFromWKB -> GEOMETRY column.
+    Skips if the duckdb spatial extension cannot be loaded (needs network on
+    first install).
+    """
+    if gpd is None:
+        pytest.skip("geopandas is not installed")
+    try:
+        source = DuckDBSource(
+            uri=':memory:',
+            initializers=["INSTALL spatial;", "LOAD spatial;"],
+            tables={'geo': 'SELECT * FROM geo_tbl'},
+        )
+    except Exception as e:  # pragma: no cover - environment dependent
+        pytest.skip(f"duckdb spatial extension unavailable: {e}")
+    gdf = gpd.GeoDataFrame(
+        {
+            'name': ['a', 'b'],
+            'pop': [1, 2],
+            'geometry': [
+                Polygon([(0, 0), (1, 0), (1, 1)]),
+                Polygon([(2, 0), (3, 0), (3, 1)]),
+            ],
+        },
+        crs='EPSG:4326',
+    )
+    tmp = pd.DataFrame(
+        {'name': gdf['name'], 'pop': gdf['pop'], 'geometry': gdf['geometry'].to_wkb()}
+    )
+    source._connection.register('geo_temp', tmp)
+    source._connection.execute(
+        'CREATE TABLE geo_tbl AS '
+        'SELECT name, pop, ST_GeomFromWKB(geometry) AS geometry FROM geo_temp'
+    )
+    return source, gpd
+
+
+def test_duckdb_geometry_returns_geodataframe():
+    """A native GEOMETRY column round-trips to a GeoDataFrame without crashing."""
+    source, gpd = _spatial_source()
+    result = source.get('geo')
+    assert isinstance(result, gpd.GeoDataFrame)
+    assert 'geometry' in result.columns
+    assert len(result) == 2
+    assert result.geometry.notna().all()
+    assert str(result.geometry.dtype) == 'geometry'
+
+
+def test_duckdb_geometry_crs_none_by_default():
+    """Without geometry_crs set the CRS stays None (WKB carries none); no regression."""
+    source, gpd = _spatial_source()
+    assert source.get('geo').crs is None
+
+
+def test_duckdb_geometry_crs_preserved():
+    """geometry_crs is reapplied when rebuilding the GeoDataFrame (gh-1904)."""
+    source, gpd = _spatial_source()
+    source.geometry_crs = 'EPSG:4326'
+    result = source.get('geo')
+    assert result.crs is not None
+    assert result.crs.to_epsg() == 4326
+
+
+def test_duckdb_geometry_crs_propagates_to_derived_source():
+    """A source created via create_sql_expr_source keeps geometry_crs (gh-1904)."""
+    source, gpd = _spatial_source()
+    source.geometry_crs = 'EPSG:4326'
+    derived = source.create_sql_expr_source(
+        {'geo2': 'SELECT * FROM geo'}, materialize=False
+    )
+    assert derived.geometry_crs == 'EPSG:4326'
+    assert derived.get('geo').crs.to_epsg() == 4326
+
+
+def test_duckdb_get_schema_geometry_no_distinct():
+    """get_schema on a geometry table returns a geospatial marker and does not
+    run DISTINCT/min-max on the geometry column."""
+    source, gpd = _spatial_source()
+    schema = source.get_schema('geo')
+    assert schema['geometry']['format'] == 'geometry'
+    assert 'enum' not in schema['geometry']
+    assert 'inclusiveMinimum' not in schema['geometry']
+    # non-geometry columns still summarised as usual
+    assert schema['pop']['inclusiveMinimum'] == 1
+    assert schema['pop']['inclusiveMaximum'] == 2

lumen/tests/sources/test_optional_imports.py

@@ -14,7 +14,7 @@
 
 import pytest
 
-from lumen.util import try_import_xarray
+from lumen.util import try_import_geopandas, try_import_xarray
 
 # Each entry: (module_path, class_name, guard_package, pip_extra)
 OPTIONAL_SOURCES = [
@@ -88,6 +88,32 @@ def mock_import(name, *args, **kwargs):
         sys.modules.update(saved)
 
 
+def test_try_import_geopandas_returns_module_when_installed():
+    """try_import_geopandas returns the geopandas module when importable."""
+    gpd = pytest.importorskip("geopandas")
+    assert try_import_geopandas() is gpd
+
+
+def test_try_import_geopandas_none_when_missing():
+    """try_import_geopandas returns None (not raises) when geopandas is absent."""
+    real_import = builtins.__import__
+
+    def mock_import(name, *args, **kwargs):
+        if name == "geopandas" or name.startswith("geopandas."):
+            raise ImportError("No module named 'geopandas'")
+        return real_import(name, *args, **kwargs)
+
+    saved = {
+        key: sys.modules.pop(key)
+        for key in [k for k in sys.modules if k == "geopandas" or k.startswith("geopandas.")]
+    }
+    try:
+        with patch("builtins.__import__", side_effect=mock_import):
+            assert try_import_geopandas() is None
+    finally:
+        sys.modules.update(saved)
+
+
 def test_try_import_xarray_returns_module_when_installed():
     """try_import_xarray returns the xarray module when xarray and xarray-sql are installed."""
     xr = pytest.importorskip("xarray")