[arrow-adbc] branch main updated: feat(python/adbc_driver_postgresql): add PostgreSQL options enum (#886)

[li...@apache.org]

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 b4b5a37a feat(python/adbc_driver_postgresql): add PostgreSQL options enum (#886)
b4b5a37a is described below

commit b4b5a37a990a3e08ca11a95d73041a3b4fcc1921
Author: David Li <li...@gmail.com>
AuthorDate: Tue Jul 11 11:05:58 2023 -0400

    feat(python/adbc_driver_postgresql): add PostgreSQL options enum (#886)
    
    Fixes #882.
---
 docs/source/python/api/adbc_driver_postgresql.rst  |  4 +--
 .../adbc_driver_postgresql/__init__.py             | 13 ++++++++-
 python/adbc_driver_postgresql/tests/test_dbapi.py  | 33 +++++++++++++++++++++-
 3 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/docs/source/python/api/adbc_driver_postgresql.rst b/docs/source/python/api/adbc_driver_postgresql.rst
index 397f979a..d7c07e91 100644
--- a/docs/source/python/api/adbc_driver_postgresql.rst
+++ b/docs/source/python/api/adbc_driver_postgresql.rst
@@ -25,8 +25,8 @@ Low-Level API
 =============
 
 .. automodule:: adbc_driver_postgresql
-
-.. autofunction:: adbc_driver_postgresql.connect
+   :members:
+   :undoc-members:
 
 DBAPI 2.0 API
 =============
diff --git a/python/adbc_driver_postgresql/adbc_driver_postgresql/__init__.py b/python/adbc_driver_postgresql/adbc_driver_postgresql/__init__.py
index 7df15045..bd900492 100644
--- a/python/adbc_driver_postgresql/adbc_driver_postgresql/__init__.py
+++ b/python/adbc_driver_postgresql/adbc_driver_postgresql/__init__.py
@@ -15,13 +15,24 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import enum
 import functools
 
 import adbc_driver_manager
 
 from ._version import __version__
 
-__all__ = ["connect", "__version__"]
+__all__ = ["StatementOptions", "connect", "__version__"]
+
+
+class StatementOptions(enum.Enum):
+    """Statement options specific to the PostgreSQL driver."""
+
+    #: Try to limit returned batches to this size (in bytes).
+    #:
+    #: This is merely a hint, and because the size is estimated, the
+    #: actual size may differ.
+    BATCH_SIZE_HINT_BYTES = "adbc.postgresql.batch_size_hint_bytes"
 
 
 def connect(uri: str) -> adbc_driver_manager.AdbcDatabase:
diff --git a/python/adbc_driver_postgresql/tests/test_dbapi.py b/python/adbc_driver_postgresql/tests/test_dbapi.py
index 8d9b09b5..c50cad1e 100644
--- a/python/adbc_driver_postgresql/tests/test_dbapi.py
+++ b/python/adbc_driver_postgresql/tests/test_dbapi.py
@@ -19,7 +19,7 @@ from typing import Generator
 
 import pytest
 
-from adbc_driver_postgresql import dbapi
+from adbc_driver_postgresql import StatementOptions, dbapi
 
 
 @pytest.fixture
@@ -28,6 +28,37 @@ def postgres(postgres_uri: str) -> Generator[dbapi.Connection, None, None]:
         yield conn
 
 
+def test_query_batch_size(postgres: dbapi.Connection):
+    with postgres.cursor() as cur:
+        cur.execute("DROP TABLE IF EXISTS test_batch_size")
+        cur.execute("CREATE TABLE test_batch_size (ints INT)")
+        cur.execute(
+            """
+            INSERT INTO test_batch_size (ints)
+            SELECT generated :: INT
+            FROM GENERATE_SERIES(1, 65536) temp(generated)
+        """
+        )
+
+        cur.execute("SELECT * FROM test_batch_size")
+        table = cur.fetch_arrow_table()
+        assert len(table.to_batches()) == 1
+
+        cur.adbc_statement.set_options(
+            **{StatementOptions.BATCH_SIZE_HINT_BYTES.value: "1"}
+        )
+        cur.execute("SELECT * FROM test_batch_size")
+        table = cur.fetch_arrow_table()
+        assert len(table.to_batches()) == 65536
+
+        cur.adbc_statement.set_options(
+            **{StatementOptions.BATCH_SIZE_HINT_BYTES.value: "4096"}
+        )
+        cur.execute("SELECT * FROM test_batch_size")
+        table = cur.fetch_arrow_table()
+        assert 64 <= len(table.to_batches()) <= 256
+
+
 def test_query_trivial(postgres: dbapi.Connection):
     with postgres.cursor() as cur:
         cur.execute("SELECT 1")