diff --git a/ibis/backends/bigquery/__init__.py b/ibis/backends/bigquery/__init__.py index 84a0d9de0000..a81958757fff 100644 --- a/ibis/backends/bigquery/__init__.py +++ b/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 = [ diff --git a/ibis/backends/bigquery/converter.py b/ibis/backends/bigquery/converter.py index f8e54afe680b..d402d1d50ec3 100644 --- a/ibis/backends/bigquery/converter.py +++ b/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 diff --git a/ibis/backends/bigquery/tests/system/test_client.py b/ibis/backends/bigquery/tests/system/test_client.py index f44893ff17fb..bc5d7eddca33 100644 --- a/ibis/backends/bigquery/tests/system/test_client.py +++ b/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 diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/difference/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/difference/out.sql index 6a35d9cfa05e..187362ed2676 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/difference/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_difference(`t0`.`geog0`, `t0`.`geog1`) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/intersection/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/intersection/out.sql index 5dfd3ab8b909..ce7d7ef65299 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/intersection/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_intersection(`t0`.`geog0`, `t0`.`geog1`) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/union/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/union/out.sql index 2d55fe586ad8..8c0653404d5c 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_binary/union/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_union(`t0`.`geog0`, `t0`.`geog1`) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_point/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_point/out.sql index a7e665427598..07a497181742 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_point/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_geogpoint(`t0`.`lon`, `t0`.`lat`) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_simplify/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_simplify/out.sql index e5d6f7d549ef..d8dbe16dd0fd 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_simplify/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_simplify(`t0`.`geog`, 5.2) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/buffer/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/buffer/out.sql index 70cb2b7cc351..d5f9e492997a 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/buffer/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_buffer(`t0`.`geog`, 5.2) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/centroid/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/centroid/out.sql index 446af3e347eb..2fde8567311f 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/centroid/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_centroid(`t0`.`geog`) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/end_point/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/end_point/out.sql index 9bb5289851e1..b9f0848346da 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/end_point/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_endpoint(`t0`.`geog`) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/point_n/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/point_n/out.sql index 070963810106..2975cc4a75bd 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/point_n/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_pointn(`t0`.`geog`, 3) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/start_point/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/start_point/out.sql index a4681f1837fd..8bd74147fadd 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary/start_point/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_startpoint(`t0`.`geog`) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file diff --git a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary_union/out.sql b/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary_union/out.sql index 4729efce3601..9d358ff759b5 100644 --- a/ibis/backends/bigquery/tests/unit/snapshots/test_compiler/test_geospatial_unary_union/out.sql +++ b/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` \ No newline at end of file + * + REPLACE (st_asbinary(`tmp`) AS `tmp`) +FROM ( + SELECT + st_union_agg(`t0`.`geog`) AS `tmp` + FROM `t` AS `t0` +) \ No newline at end of file