From f89aa32e8de1f4a7d2bcc7379c690003fe4e2b51 Mon Sep 17 00:00:00 2001 From: Phillip Cloud <417981+cpcloud@users.noreply.github.com> Date: Mon, 29 Jul 2024 15:04:54 -0400 Subject: [PATCH] refactor(duckdb): use replace to generate less sql (#9713) Refactors the duckdb output to geospatial to use `* REPLACE` syntax to generate less code when replacing a small number of geospatial columns (usually one) among an otherwise large set of columns. Additionally, we now only need to loop over expression columns once. --- ibis/backends/duckdb/__init__.py | 20 ++++++++++--------- .../expr0/out.sql | 3 ++- .../expr1/out.sql | 3 ++- .../shp0-0_0/out.sql | 3 ++- .../shp1-1_1/out.sql | 3 ++- .../shp2-2_2/out.sql | 3 ++- .../shp3-0_0_1_1_2_2/out.sql | 3 ++- .../shp4-2_2_1_1_0_0/out.sql | 3 ++- .../shp5-0_0_1_1_2_2_0_0/out.sql | 3 ++- .../shp6-0_0_1_1_2_2_0_0/out.sql | 3 ++- .../shp7-0_0_1_1_2_2_2_2_1_1_0_0/out.sql | 3 ++- 11 files changed, 31 insertions(+), 19 deletions(-) diff --git a/ibis/backends/duckdb/__init__.py b/ibis/backends/duckdb/__init__.py index 28eb75e14340..4914f8a0b993 100644 --- a/ibis/backends/duckdb/__init__.py +++ b/ibis/backends/duckdb/__init__.py @@ -101,23 +101,25 @@ def _to_sqlglot( sql = super()._to_sqlglot(expr, limit=limit, params=params) table_expr = expr.as_table() - geocols = frozenset( + geocols = [ name for name, typ in table_expr.schema().items() if typ.is_geospatial() - ) + ] if not geocols: return sql else: self._load_extensions(["spatial"]) + compiler = self.compiler + quoted = compiler.quoted return sg.select( - *( - self.compiler.f.st_aswkb( - sg.column(col, quoted=self.compiler.quoted) - ).as_(col) - if col in geocols - else col - for col in table_expr.columns + sge.Star( + replace=[ + compiler.f.st_aswkb(sg.column(col, quoted=quoted)).as_( + col, quoted=quoted + ) + for col in geocols + ] ) ).from_(sql.subquery()) diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_explicit/expr0/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_explicit/expr0/out.sql index 72eb287d8221..29ac95d4fc37 100644 --- a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_explicit/expr0/out.sql +++ b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_explicit/expr0/out.sql @@ -1,5 +1,6 @@ SELECT - ST_ASWKB("p") AS p + * + REPLACE (ST_ASWKB("p") AS "p") FROM ( SELECT ST_GEOMFROMTEXT('POINT (1 0)') AS "p" diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_explicit/expr1/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_explicit/expr1/out.sql index 72eb287d8221..29ac95d4fc37 100644 --- a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_explicit/expr1/out.sql +++ b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_explicit/expr1/out.sql @@ -1,5 +1,6 @@ SELECT - ST_ASWKB("p") AS p + * + REPLACE (ST_ASWKB("p") AS "p") FROM ( SELECT ST_GEOMFROMTEXT('POINT (1 0)') AS "p" diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp0-0_0/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp0-0_0/out.sql index 58f29445ffed..2bab8ff0b496 100644 --- a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp0-0_0/out.sql +++ b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp0-0_0/out.sql @@ -1,5 +1,6 @@ SELECT - ST_ASWKB("result") AS result + * + REPLACE (ST_ASWKB("result") AS "result") FROM ( SELECT ST_GEOMFROMTEXT('POINT (0 0)') AS "result" diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp1-1_1/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp1-1_1/out.sql index 7fabbfbb4cee..007a60962579 100644 --- a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp1-1_1/out.sql +++ b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp1-1_1/out.sql @@ -1,5 +1,6 @@ SELECT - ST_ASWKB("result") AS result + * + REPLACE (ST_ASWKB("result") AS "result") FROM ( SELECT ST_GEOMFROMTEXT('POINT (1 1)') AS "result" diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp2-2_2/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp2-2_2/out.sql index 203b31c741bb..61a580eefa12 100644 --- a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp2-2_2/out.sql +++ b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp2-2_2/out.sql @@ -1,5 +1,6 @@ SELECT - ST_ASWKB("result") AS result + * + REPLACE (ST_ASWKB("result") AS "result") FROM ( SELECT ST_GEOMFROMTEXT('POINT (2 2)') AS "result" diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp3-0_0_1_1_2_2/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp3-0_0_1_1_2_2/out.sql index a825f7c68d8f..e4594f6892df 100644 --- a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp3-0_0_1_1_2_2/out.sql +++ b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp3-0_0_1_1_2_2/out.sql @@ -1,5 +1,6 @@ SELECT - ST_ASWKB("result") AS result + * + REPLACE (ST_ASWKB("result") AS "result") FROM ( SELECT ST_GEOMFROMTEXT('LINESTRING (0 0, 1 1, 2 2)') AS "result" diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp4-2_2_1_1_0_0/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp4-2_2_1_1_0_0/out.sql index 193852e48ca6..a4cb44ad3495 100644 --- a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp4-2_2_1_1_0_0/out.sql +++ b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp4-2_2_1_1_0_0/out.sql @@ -1,5 +1,6 @@ SELECT - ST_ASWKB("result") AS result + * + REPLACE (ST_ASWKB("result") AS "result") FROM ( SELECT ST_GEOMFROMTEXT('LINESTRING (2 2, 1 1, 0 0)') AS "result" diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp5-0_0_1_1_2_2_0_0/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp5-0_0_1_1_2_2_0_0/out.sql index a06759136ca8..3240c2f13c10 100644 --- a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp5-0_0_1_1_2_2_0_0/out.sql +++ b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp5-0_0_1_1_2_2_0_0/out.sql @@ -1,5 +1,6 @@ SELECT - ST_ASWKB("result") AS result + * + REPLACE (ST_ASWKB("result") AS "result") FROM ( SELECT ST_GEOMFROMTEXT('POLYGON ((0 0, 1 1, 2 2, 0 0))') AS "result" diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp6-0_0_1_1_2_2_0_0/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp6-0_0_1_1_2_2_0_0/out.sql index 3631595fc330..7e6330703a19 100644 --- a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp6-0_0_1_1_2_2_0_0/out.sql +++ b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp6-0_0_1_1_2_2_0_0/out.sql @@ -1,5 +1,6 @@ SELECT - ST_ASWKB("result") AS result + * + REPLACE (ST_ASWKB("result") AS "result") FROM ( SELECT ST_GEOMFROMTEXT('MULTIPOLYGON (((0 0, 1 1, 2 2, 0 0)))') AS "result" diff --git a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp7-0_0_1_1_2_2_2_2_1_1_0_0/out.sql b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp7-0_0_1_1_2_2_2_2_1_1_0_0/out.sql index a364a2d0e48e..1f9d59c11334 100644 --- a/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp7-0_0_1_1_2_2_2_2_1_1_0_0/out.sql +++ b/ibis/backends/duckdb/tests/snapshots/test_geospatial/test_literal_geospatial_inferred/shp7-0_0_1_1_2_2_2_2_1_1_0_0/out.sql @@ -1,5 +1,6 @@ SELECT - ST_ASWKB("result") AS result + * + REPLACE (ST_ASWKB("result") AS "result") FROM ( SELECT ST_GEOMFROMTEXT('MULTILINESTRING ((0 0, 1 1, 2 2), (2 2, 1 1, 0 0))') AS "result"