python tests

.gitignore

@@ -36,4 +36,5 @@
 build
 tmp
 .vscode
-duckdb_unittest_tempdir
+duckdb_unittest_tempdir
+__pycache__

test/python/README.md

@@ -0,0 +1,18 @@
+
+# Running Python tests
+
+First, a development build of the duckdb Python package that matches the commit
+of the duckdb submodule used to build the extension.
+
+```shell
+cd duckdb/tools/pythonpkg
+pip install -r requirements-dev.txt
+pip install .
+```
+
+Running the `test_*.py` files requires `pytest`:
+
+```shell
+# pip install pytest pyarrow
+pytest -vv tests/python
+```

test/python/conftest.py

@@ -0,0 +1,40 @@
+import glob
+from pathlib import Path
+
+import pytest
+import duckdb
+import warnings
+
+
+HERE = Path(__file__).parent
+
+
+def _install_dev_and_connect():
+    con = duckdb.connect(config={"allow_unsigned_extensions": True})
+
+    possible_builds = glob.glob(
+        "build/**/spatial/spatial.duckdb_extension",
+        recursive=True,
+        root_dir=HERE.parent.parent,
+    )
+    if possible_builds:
+        con.install_extension(possible_builds[0], force_install=True)
+    else:
+        warnings.warn(
+            "Can't find build directory for spatial.duckdb_extension; skipping INSTALL"
+        )
+
+    con.load_extension("spatial")
+    return con
+
+
+@pytest.fixture()
+def geoarrow_con():
+    con = _install_dev_and_connect()
+    con.sql("""CALL register_geoarrow_extensions()""")
+    return con
+
+
+@pytest.fixture()
+def con():
+    return _install_dev_and_connect()

test/python/test_geoarrow.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+
+import pyarrow as pa
+from pyarrow import parquet
+
+HERE = Path(__file__).parent
+
+
+def test_export_without_register(con):
+    tab = con.sql("""SELECT ST_GeomFromText('POINT (0 1)') as geom;""").to_arrow_table()
+    assert tab.schema.field("geom").metadata is None
+
+
+def test_basic_export(geoarrow_con):
+    tab = geoarrow_con.sql(
+        """SELECT ST_GeomFromText('POINT (0 1)') as geom;"""
+    ).to_arrow_table()
+    assert tab.schema.field("geom").metadata == {
+        b"ARROW:extension:metadata": b"{}",
+        b"ARROW:extension:name": b"geoarrow.wkb",
+    }
+
+
+def test_basic_import(geoarrow_con):
+    tab = geoarrow_con.sql(
+        """SELECT ST_GeomFromText('POINT (0 1)') as geom;"""
+    ).to_arrow_table()
+    assert tab.schema.field("geom").metadata == {
+        b"ARROW:extension:metadata": b"{}",
+        b"ARROW:extension:name": b"geoarrow.wkb",
+    }
+
+
+def test_roundtrip_segments(geoarrow_con):
+    segments_file = HERE.parent / "data" / "segments.parquet"
+    raw_table = parquet.read_table(segments_file)
+    assert raw_table.schema.field("geometry").metadata is None
+
+    field = pa.field(
+        "geometry", pa.binary(), metadata={"ARROW:extension:name": "geoarrow.wkb"}
+    )
+    schema = pa.schema([field])
+    geo_table = pa.table([raw_table["geometry"]], schema=schema)
+
+    geoarrow_table_wkt = geoarrow_con.sql(
+        """SELECT ST_AsText(geometry) as wkt FROM geo_table"""
+    ).to_arrow_table()
+    geoparquet_table_wkt = geoarrow_con.sql(
+        f"""SELECT ST_AsText(geometry) as wkt FROM "{segments_file}";""",
+    ).to_arrow_table()
+    assert geoarrow_table_wkt == geoparquet_table_wkt
+
+    # While we're here, check the roundtrip output
+    assert geoarrow_con.sql("""SELECT * from geo_table""").to_arrow_table() == geo_table