feat(api): add to_geo methods for writing geospatial output

ibis-project · Oct 10, 2024 · bdc3484 · bdc3484
1 parent 5be8232
commit bdc3484
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 0 deletions.
diff --git 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
@@ -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)