This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new 2b8e429 test(python/adbc_driver_postgres): add polars integration
test (#542)
2b8e429 is described below
commit 2b8e429f90fe730f7cb1df769d7f69f1e43ac117
Author: David Li <[email protected]>
AuthorDate: Thu Mar 30 11:39:37 2023 -0400
test(python/adbc_driver_postgres): add polars integration test (#542)
Fixes #539.
---
ci/conda_env_python.txt | 3 +
ci/scripts/python_util.sh | 7 +-
ci/scripts/python_wheel_windows_test.bat | 2 +-
python/adbc_driver_postgresql/pyproject.toml | 7 +-
.../tests/{test_dbapi.py => conftest.py} | 31 +++++---
python/adbc_driver_postgresql/tests/test_dbapi.py | 9 +--
.../adbc_driver_postgresql/tests/test_lowlevel.py | 7 +-
python/adbc_driver_postgresql/tests/test_polars.py | 91 ++++++++++++++++++++++
8 files changed, 127 insertions(+), 30 deletions(-)
diff --git a/ci/conda_env_python.txt b/ci/conda_env_python.txt
index 4c8b129..46690ee 100644
--- a/ci/conda_env_python.txt
+++ b/ci/conda_env_python.txt
@@ -20,3 +20,6 @@ pandas
pyarrow>=8.0.0
pytest
setuptools
+
+# For integration testing
+polars
diff --git a/ci/scripts/python_util.sh b/ci/scripts/python_util.sh
index ad67612..b03af24 100644
--- a/ci/scripts/python_util.sh
+++ b/ci/scripts/python_util.sh
@@ -171,9 +171,10 @@ import $component.dbapi
# --import-mode required, else tries to import from the source dir
instead of installed package
if [[ "${component}" = "adbc_driver_manager" ]]; then
- python -m pytest -vvx --import-mode append -k "not sqlite"
${source_dir}/python/$component/tests
- else
- python -m pytest -vvx --import-mode append
${source_dir}/python/$component/tests
+ export PYTEST_ADDOPTS="-k 'not sqlite'"
+ elif [[ "${component}" = "adbc_driver_postgresql" ]]; then
+ export PYTEST_ADDOPTS="-k 'not polars'"
fi
+ python -m pytest -vvx --import-mode append
${source_dir}/python/$component/tests
done
}
diff --git a/ci/scripts/python_wheel_windows_test.bat
b/ci/scripts/python_wheel_windows_test.bat
index f142705..6459529 100644
--- a/ci/scripts/python_wheel_windows_test.bat
+++ b/ci/scripts/python_wheel_windows_test.bat
@@ -35,5 +35,5 @@ FOR %%c IN (adbc_driver_manager adbc_driver_flightsql
adbc_driver_postgresql adb
echo "=== Testing %%c ==="
python -c "import %%c" || exit /B 1
python -c "import %%c.dbapi" || exit /B 1
- python -m pytest -vvx --import-mode append -k "not sqlite"
%source_dir%\python\%%c\tests || exit /B 1
+ python -m pytest -vvx --import-mode append -k "not sqlite and not polars"
%source_dir%\python\%%c\tests || exit /B 1
)
diff --git a/python/adbc_driver_postgresql/pyproject.toml
b/python/adbc_driver_postgresql/pyproject.toml
index 48deea2..bc7da31 100644
--- a/python/adbc_driver_postgresql/pyproject.toml
+++ b/python/adbc_driver_postgresql/pyproject.toml
@@ -28,7 +28,7 @@ dependencies = [
[project.optional-dependencies]
dbapi = ["pandas", "pyarrow>=8.0.0"]
-test = ["pandas", "pyarrow>=8.0.0", "pytest"]
+test = ["pandas", "polars", "pyarrow>=8.0.0", "pytest"]
[project.urls]
homepage = "https://arrow.apache.org";
@@ -38,6 +38,11 @@ repository = "https://github.com/apache/arrow-adbc";
requires = ["setuptools >= 61.0.0"]
build-backend = "setuptools.build_meta"
+[tool.pytest.ini_options]
+markers = [
+ "polars: integration tests with polars",
+]
+
[tool.setuptools]
include-package-data = true
license-files = ["LICENSE.txt", "NOTICE.txt"]
diff --git a/python/adbc_driver_postgresql/tests/test_dbapi.py
b/python/adbc_driver_postgresql/tests/conftest.py
similarity index 52%
copy from python/adbc_driver_postgresql/tests/test_dbapi.py
copy to python/adbc_driver_postgresql/tests/conftest.py
index c3223e8..e651781 100644
--- a/python/adbc_driver_postgresql/tests/test_dbapi.py
+++ b/python/adbc_driver_postgresql/tests/conftest.py
@@ -14,24 +14,31 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
import os
import pytest
-from adbc_driver_postgresql import dbapi
-
[email protected]
-def postgres():
[email protected](scope="session")
+def postgres_uri() -> str:
postgres_uri = os.environ.get("ADBC_POSTGRESQL_TEST_URI")
if not postgres_uri:
pytest.skip("Set ADBC_POSTGRESQL_TEST_URI to run tests")
- with dbapi.connect(postgres_uri) as conn:
- yield conn
-
-
-def test_query_trivial(postgres):
- with postgres.cursor() as cur:
- cur.execute("SELECT 1")
- assert cur.fetchone() == (1,)
+ return postgres_uri
diff --git a/python/adbc_driver_postgresql/tests/test_dbapi.py
b/python/adbc_driver_postgresql/tests/test_dbapi.py
index c3223e8..7344f1a 100644
--- a/python/adbc_driver_postgresql/tests/test_dbapi.py
+++ b/python/adbc_driver_postgresql/tests/test_dbapi.py
@@ -15,23 +15,18 @@
# specific language governing permissions and limitations
# under the License.
-import os
-
import pytest
from adbc_driver_postgresql import dbapi
@pytest.fixture
-def postgres():
- postgres_uri = os.environ.get("ADBC_POSTGRESQL_TEST_URI")
- if not postgres_uri:
- pytest.skip("Set ADBC_POSTGRESQL_TEST_URI to run tests")
+def postgres(postgres_uri: str) -> dbapi.Connection:
with dbapi.connect(postgres_uri) as conn:
yield conn
-def test_query_trivial(postgres):
+def test_query_trivial(postgres: dbapi.Connection):
with postgres.cursor() as cur:
cur.execute("SELECT 1")
assert cur.fetchone() == (1,)
diff --git a/python/adbc_driver_postgresql/tests/test_lowlevel.py
b/python/adbc_driver_postgresql/tests/test_lowlevel.py
index e5c44f5..587f7b8 100644
--- a/python/adbc_driver_postgresql/tests/test_lowlevel.py
+++ b/python/adbc_driver_postgresql/tests/test_lowlevel.py
@@ -15,8 +15,6 @@
# specific language governing permissions and limitations
# under the License.
-import os
-
import pyarrow
import pytest
@@ -25,10 +23,7 @@ import adbc_driver_postgresql
@pytest.fixture
-def postgres():
- postgres_uri = os.environ.get("ADBC_POSTGRESQL_TEST_URI")
- if not postgres_uri:
- pytest.skip("Set ADBC_POSTGRESQL_TEST_URI to run tests")
+def postgres(postgres_uri: str) -> adbc_driver_manager.AdbcConnection:
with adbc_driver_postgresql.connect(postgres_uri) as db:
with adbc_driver_manager.AdbcConnection(db) as conn:
yield conn
diff --git a/python/adbc_driver_postgresql/tests/test_polars.py
b/python/adbc_driver_postgresql/tests/test_polars.py
new file mode 100644
index 0000000..d03bf53
--- /dev/null
+++ b/python/adbc_driver_postgresql/tests/test_polars.py
@@ -0,0 +1,91 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Integration tests with polars."""
+
+import uuid
+
+import pytest
+
+from adbc_driver_postgresql import dbapi
+
+try:
+ import polars
+except ImportError:
+ pass
+
+pytestmark = pytest.mark.polars
+
+
[email protected]
+def df(request):
+ if request.param == "ints":
+ return polars.DataFrame(
+ {
+ "ints": [1, 2, 4, 8],
+ },
+ schema={
+ "ints": polars.Int64,
+ },
+ )
+ elif request.param == "floats":
+ return polars.DataFrame(
+ {
+ "floats": [1, 2, 4, 8],
+ },
+ schema={
+ "floats": polars.Float64,
+ },
+ )
+ raise KeyError(f"Unknown df {request.param}")
+
+
[email protected](
+ "df",
+ [
+ "ints",
+ pytest.param(
+ "floats",
+ marks=pytest.mark.xfail(
+ reason="; ".join(
+ [
+ "apache/arrow-adbc#81: lack of type support",
+ "pola-rs/polars#7757: polars doesn't close cursor
properly",
+ ]
+ )
+ ),
+ ),
+ ],
+ indirect=True,
+)
+def test_polars_write_database(postgres_uri: str, df: "polars.DataFrame") ->
None:
+ table_name = f"polars_test_ingest_{uuid.uuid4().hex}"
+ try:
+ df.write_database(
+ table_name=table_name,
+ connection_uri=postgres_uri,
+ # TODO(apache/arrow-adbc#541): polars doesn't map the semantics
+ # properly here, and one of their modes isn't supported
+ if_exists="replace",
+ engine="adbc",
+ )
+ finally:
+ # TODO(apache/arrow-adbc#540): driver doesn't handle execute()
+ # properly here because it tries to infer the schema.
+ with dbapi.connect(postgres_uri) as conn:
+ with conn.cursor() as cursor:
+ cursor.executemany(f"DROP TABLE {table_name}", [])
