Skip to content

Commit

Permalink
fix(bigquery): repr geospatial values in interactive mode (#9712)
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud authored Jul 30, 2024
1 parent 9751a36 commit bd8c93f
Show file tree
Hide file tree
Showing 14 changed files with 117 additions and 24 deletions.
26 changes: 25 additions & 1 deletion ibis/backends/bigquery/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down
2 changes: 1 addition & 1 deletion ibis/backends/bigquery/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 14 additions & 0 deletions ibis/backends/bigquery/tests/system/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Original file line number Diff line number Diff line change
@@ -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`
)
Original file line number Diff line number Diff line change
@@ -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`
)
Original file line number Diff line number Diff line change
@@ -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`
)
Original file line number Diff line number Diff line change
@@ -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`
)
Original file line number Diff line number Diff line change
@@ -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`
)
Original file line number Diff line number Diff line change
@@ -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`
)
Original file line number Diff line number Diff line change
@@ -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`
)
Original file line number Diff line number Diff line change
@@ -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`
)
Original file line number Diff line number Diff line change
@@ -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`
)
Original file line number Diff line number Diff line change
@@ -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`
)
Original file line number Diff line number Diff line change
@@ -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`
)

0 comments on commit bd8c93f

Please sign in to comment.