diff --git a/ibis/backends/duckdb/__init__.py b/ibis/backends/duckdb/__init__.py index 210670164316..5219d249e305 100644 --- a/ibis/backends/duckdb/__init__.py +++ b/ibis/backends/duckdb/__init__.py @@ -493,6 +493,49 @@ def read_csv( con.exec_driver_sql(view) return self.table(table_name) + def read_geo( + self, + source: str, + table_name: str | None = None, + **kwargs: Any, + ) -> ir.Table: + """Register a GEO file as a table in the current database. + + Parameters + ---------- + source + The data source(s). Path to a file of geospatial files supported + by duckdb. + See https://duckdb.org/docs/extensions/spatial.html#st_read---read-spatial-data-from-files + table_name + An optional name to use for the created table. This defaults to + a sequentially generated name. + **kwargs + Additional keyword arguments passed to DuckDB loading function. + See https://duckdb.org/docs/extensions/spatial.html#st_read---read-spatial-data-from-files + for more information. + + Returns + ------- + ir.Table + The just-registered table + """ + + if not table_name: + table_name = util.gen_name("read_geo") + + # load geospatial extension + self.load_extension("spatial") + + source_expr = sa.select(sa.literal_column("*")).select_from( + sa.func.st_read(util.normalize_filename(source), _format_kwargs(kwargs)) + ) + + view = self._compile_temp_view(table_name, source_expr) + with self.begin() as con: + con.exec_driver_sql(view) + return self.table(table_name) + def read_parquet( self, source_list: str | Iterable[str], diff --git a/ibis/backends/duckdb/tests/conftest.py b/ibis/backends/duckdb/tests/conftest.py index 2c96dd8ca852..d873591904bf 100644 --- a/ibis/backends/duckdb/tests/conftest.py +++ b/ibis/backends/duckdb/tests/conftest.py @@ -81,6 +81,13 @@ def con(data_dir, tmp_path_factory, worker_id): return TestConf.load_data(data_dir, tmp_path_factory, worker_id).connection +@pytest.fixture(scope="session") +def zones(con, data_dir): + # pending merge https://github.com/ibis-project/testing-data/pull/5 + zones = con.read_geo(data_dir / "geojson" / "zones.geojson") + return zones + + @pytest.fixture(scope="session") def zones_gdf(data_dir): # pending merge https://github.com/ibis-project/testing-data/pull/5 diff --git a/ibis/backends/duckdb/tests/test_register.py b/ibis/backends/duckdb/tests/test_register.py index 50c1682e4264..8cc913946e61 100644 --- a/ibis/backends/duckdb/tests/test_register.py +++ b/ibis/backends/duckdb/tests/test_register.py @@ -46,6 +46,18 @@ def test_read_parquet(data_dir): assert t.count().execute() +@pytest.mark.xfail(raises=duckdb.duckdb.CatalogException, reason="ST_AsEWKB") +def test_read_geo_fail(con, data_dir): + t = con.read_geo(data_dir / "geojson" / "zones.geojson") + # can't convert geometry to arrow type yet + assert t.head().to_pyarrow() + + +def test_read_geo(con, data_dir): + t = con.read_geo(data_dir / "geojson" / "zones.geojson") + assert t.count().execute() + + @pytest.mark.xfail_version( duckdb=["duckdb<0.7.0"], reason="read_json_auto doesn't exist", raises=exc.IbisError )