fix(bigquery): repr geospatial values in interactive mode (#9712)

Author: cpcloud

ibis/backends/bigquery/__init__.py

@@ -683,12 +683,36 @@ def _to_sqlglot(
         self._define_udf_translation_rules(expr)
         sql = super()._to_sqlglot(expr, limit=limit, params=params, **kwargs)
 
+        table_expr = expr.as_table()
+        geocols = [
+            name for name, typ in table_expr.schema().items() if typ.is_geospatial()
+        ]
+
         query = sql.transform(
             _qualify_memtable,
             dataset=getattr(self._session_dataset, "dataset_id", None),
             project=getattr(self._session_dataset, "project", None),
         ).transform(_remove_null_ordering_from_unsupported_window)
-        return query
+
+        if not geocols:
+            return query
+
+        # if there are any geospatial columns, we have to convert them to WKB,
+        # so interactive mode knows how to display them
+        #
+        # by default bigquery returns data to python as WKT, and there's really
+        # no point in supporting both if we don't need to.
+        compiler = self.compiler
+        quoted = compiler.quoted
+        f = compiler.f
+        return sg.select(
+            sge.Star(
+                replace=[
+                    f.st_asbinary(sg.column(col, quoted=quoted)).as_(col, quoted=quoted)
+                    for col in geocols
+                ]
+            )
+        ).from_(query.subquery())
 
     def raw_sql(self, query: str, params=None, page_size: int | None = None):
         query_parameters = [

ibis/backends/bigquery/converter.py

@@ -9,7 +9,7 @@ def convert_GeoSpatial(cls, s, dtype, pandas_type):
         import geopandas as gpd
         import shapely as shp
 
-        return gpd.GeoSeries(shp.from_wkt(s))
+        return gpd.GeoSeries(shp.from_wkb(s))
 
     convert_Point = convert_LineString = convert_Polygon = convert_MultiLineString = (
         convert_MultiPoint

ibis/backends/bigquery/tests/system/test_client.py

@@ -455,3 +455,17 @@ def test_complex_column_name(con):
     )
     result = con.to_pandas(expr)
     assert result == 1
+
+
+def test_geospatial_interactive(con, monkeypatch):
+    pytest.importorskip("geopandas")
+
+    monkeypatch.setattr(ibis.options, "interactive", True)
+    t = con.table("bigquery-public-data.geo_us_boundaries.zip_codes")
+    expr = (
+        t.filter(lambda t: t.zip_code_geom.geometry_type() == "ST_Polygon")
+        .head(1)
+        .zip_code_geom
+    )
+    result = repr(expr)
+    assert "POLYGON" in result

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/difference/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_difference(`t0`.`geog0`, `t0`.`geog1`) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_difference(`t0`.`geog0`, `t0`.`geog1`) AS `tmp`
+  FROM `t` AS `t0`
+)

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/intersection/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_intersection(`t0`.`geog0`, `t0`.`geog1`) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_intersection(`t0`.`geog0`, `t0`.`geog1`) AS `tmp`
+  FROM `t` AS `t0`
+)

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/union/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_union(`t0`.`geog0`, `t0`.`geog1`) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_union(`t0`.`geog0`, `t0`.`geog1`) AS `tmp`
+  FROM `t` AS `t0`
+)

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_point/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_geogpoint(`t0`.`lon`, `t0`.`lat`) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_geogpoint(`t0`.`lon`, `t0`.`lat`) AS `tmp`
+  FROM `t` AS `t0`
+)

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_simplify/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_simplify(`t0`.`geog`, 5.2) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_simplify(`t0`.`geog`, 5.2) AS `tmp`
+  FROM `t` AS `t0`
+)

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/buffer/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_buffer(`t0`.`geog`, 5.2) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_buffer(`t0`.`geog`, 5.2) AS `tmp`
+  FROM `t` AS `t0`
+)

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/centroid/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_centroid(`t0`.`geog`) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_centroid(`t0`.`geog`) AS `tmp`
+  FROM `t` AS `t0`
+)

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/end_point/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_endpoint(`t0`.`geog`) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_endpoint(`t0`.`geog`) AS `tmp`
+  FROM `t` AS `t0`
+)

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/point_n/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_pointn(`t0`.`geog`, 3) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_pointn(`t0`.`geog`, 3) AS `tmp`
+  FROM `t` AS `t0`
+)

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/start_point/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_startpoint(`t0`.`geog`) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_startpoint(`t0`.`geog`) AS `tmp`
+  FROM `t` AS `t0`
+)

ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary_union/out.sql

@@ -1,3 +1,8 @@
 SELECT
-  st_union_agg(`t0`.`geog`) AS `tmp`
-FROM `t` AS `t0`
+  *
+  REPLACE (st_asbinary(`tmp`) AS `tmp`)
+FROM (
+  SELECT
+    st_union_agg(`t0`.`geog`) AS `tmp`
+  FROM `t` AS `t0`
+)