From 1235a46218c02e698d61a49b954c34832b78dcf3 Mon Sep 17 00:00:00 2001 From: Phillip Cloud <417981+cpcloud@users.noreply.github.com> Date: Thu, 10 Oct 2024 09:53:18 -0400 Subject: [PATCH] feat(api): add `to_geo` methods for writing geospatial output --- ibis/backends/duckdb/__init__.py | 29 ++++++++++ ibis/backends/duckdb/tests/test_geospatial.py | 54 +++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/ibis/backends/duckdb/__init__.py b/ibis/backends/duckdb/__init__.py index d026ab6e386b8..15a39d7c240b9 100644 --- a/ibis/backends/duckdb/__init__.py +++ b/ibis/backends/duckdb/__init__.py @@ -1583,6 +1583,35 @@ def to_csv( with self._safe_raw_sql(copy_cmd): pass + @util.experimental + def to_geo( + self, + expr: ir.Table, + path: str | Path, + *, + format: str, + params: Mapping[ir.Scalar, Any] | None = None, + limit: int | str | None = None, + layer_creation_options: Mapping[str, str] | None = None, + srs: str | int | None = None, + ) -> None: + self._run_pre_execute_hooks(expr) + query = self.compile(expr, params=params, limit=limit) + + args = [ + "FORMAT GDAL", + f"DRIVER '{format}'", + *(f"{k.upper()} {v!r}" for k, v in (layer_creation_options or {}).items()), + ] + + if srs is not None: + args.append(f"SRS {srs!r}") + + copy_cmd = f"COPY ({query}) TO {str(path)!r} ({', '.join(args)})" + + with self._safe_raw_sql(copy_cmd): + pass + def _get_schema_using_query(self, query: str) -> sch.Schema: with self._safe_raw_sql(f"DESCRIBE {query}") as cur: rows = cur.fetch_arrow_table() diff --git a/ibis/backends/duckdb/tests/test_geospatial.py b/ibis/backends/duckdb/tests/test_geospatial.py index 9ac0fea850d36..d97eb04fd5f2c 100644 --- a/ibis/backends/duckdb/tests/test_geospatial.py +++ b/ibis/backends/duckdb/tests/test_geospatial.py @@ -398,3 +398,57 @@ def test_geom_from_string(con): expr = value.cast("geometry") result = con.execute(expr) assert result == shapely.from_wkt("POINT (1 2)") + + +no_roundtrip = pytest.mark.xfail( + (duckdb.IOException, duckdb.NotImplementedException, duckdb.PermissionException), + reason="format cannot be round-tripped", +) + + +@pytest.mark.parametrize( + "driver", + [ + param("ESRI Shapefile", marks=no_roundtrip), + param("MapInfo File", marks=no_roundtrip), + param("S57", marks=no_roundtrip), + param("DGN", marks=no_roundtrip), + param("Memory", marks=no_roundtrip), + param("CSV", marks=no_roundtrip), + "GML", + param("GPX", marks=no_roundtrip), + param("KML", marks=no_roundtrip), + "GeoJSON", + "GeoJSONSeq", + param("OGR_GMT", marks=no_roundtrip), + "GPKG", + "SQLite", + param("WAsP", marks=no_roundtrip), + param("OpenFileGDB", marks=no_roundtrip), + param("DXF", marks=no_roundtrip), + param("FlatGeobuf", marks=no_roundtrip), + param("Geoconcept", marks=no_roundtrip), + param("GeoRSS", marks=no_roundtrip), + param("PGDUMP", marks=no_roundtrip), + param("GPSBabel", marks=no_roundtrip), + param("ODS", marks=no_roundtrip), + param("XLSX", marks=no_roundtrip), + param("Elasticsearch", marks=no_roundtrip), + param("Carto", marks=no_roundtrip), + param("AmigoCloud", marks=no_roundtrip), + param("Selafin", marks=no_roundtrip), + "JML", + param("VDV", marks=no_roundtrip), + param("MVT", marks=no_roundtrip), + param("NGW", marks=no_roundtrip), + "MapML", + "PMTiles", + "JSONFG", + ], + ids=lambda id: id.replace(" ", "_").lower(), +) +def test_to_geo(con, driver, zones, tmp_path): + ext = driver.replace(" ", "_").lower() + out = tmp_path / f"outfile.{ext}" + con.to_geo(zones, path=out, format=driver) + con.read_geo(out)