Skip to content

Commit

Permalink
feat(duckdb): add read_geo function
Browse files Browse the repository at this point in the history
  • Loading branch information
ncclementi authored and gforsyth committed Dec 4, 2023
1 parent 65f496c commit b19a8ce
Show file tree
Hide file tree
Showing 3 changed files with 62 additions and 0 deletions.
43 changes: 43 additions & 0 deletions ibis/backends/duckdb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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],
Expand Down
7 changes: 7 additions & 0 deletions ibis/backends/duckdb/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 12 additions & 0 deletions ibis/backends/duckdb/tests/test_register.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down

0 comments on commit b19a8ce

Please sign in to comment.