perf(postgres): improve to_pyarrow_batches by using server-side cursors (#10954)

This adds a specific `to_pyarrow_batches` implementation to the
PostgreSQL backend, which uses server side cursors. This allows ibis to
allocate memory needed only for `chunk_size` results of the query
instead of the whole set.

Resolves #10938

Author: ronif

ibis/backends/conftest.py

@@ -5,7 +5,6 @@
 import importlib.metadata
 import itertools
 from functools import cache
-from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import _pytest
@@ -148,20 +147,6 @@
 ALL_BACKENDS = set(_get_backend_names())
 
 
-@pytest.fixture(scope="session")
-def data_dir() -> Path:
-    """Return the test data directory.
-
-    Returns
-    -------
-    Path
-        Test data directory
-    """
-    root = Path(__file__).absolute().parents[2]
-
-    return root / "ci" / "ibis-testing-data"
-
-
 def _get_backend_conf(backend_str: str):
     """Convert a backend string to the test class for the backend."""
     conftest = importlib.import_module(f"ibis.backends.{backend_str}.tests.conftest")

ibis/backends/postgres/__init__.py

@@ -28,7 +28,7 @@
 from ibis.backends.sql.compilers.base import TRUE, C, ColGen
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Callable, Mapping
     from urllib.parse import ParseResult
 
     import pandas as pd
@@ -740,3 +740,53 @@ def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
         else:
             con.commit()
             return cursor
+
+    @util.experimental
+    def to_pyarrow_batches(
+        self,
+        expr: ir.Expr,
+        /,
+        *,
+        params: Mapping[ir.Scalar, Any] | None = None,
+        limit: int | str | None = None,
+        chunk_size: int = 1_000_000,
+        **_: Any,
+    ) -> pa.ipc.RecordBatchReader:
+        import pandas as pd
+        import pyarrow as pa
+
+        self._run_pre_execute_hooks(expr)
+
+        schema = expr.as_table().schema()
+
+        query = self.compile(expr, limit=limit, params=params)
+        with contextlib.suppress(AttributeError):
+            query = query.sql(dialect=self.dialect)
+
+        con = self.con
+        # server-side cursors need to be uniquely named
+        cursor = con.cursor(name=util.gen_name("postgres_cursor"))
+
+        try:
+            cursor.execute(query)
+        except Exception:
+            con.rollback()
+            cursor.close()
+            raise
+
+        def _batches(schema: pa.Schema):
+            columns = schema.names
+            try:
+                while batch := cursor.fetchmany(chunk_size):
+                    yield pa.RecordBatch.from_pandas(
+                        pd.DataFrame(batch, columns=columns), schema=schema
+                    )
+            except Exception:
+                con.rollback()
+                cursor.close()
+                raise
+            else:
+                con.commit()
+
+        pa_schema = schema.to_pyarrow()
+        return pa.RecordBatchReader.from_batches(pa_schema, _batches(pa_schema))

ibis/conftest.py

@@ -3,6 +3,7 @@
 import builtins
 import os
 import platform
+from pathlib import Path
 
 import pytest
 
@@ -44,3 +45,17 @@ def add_ibis(monkeypatch, doctest_namespace):
     condition=WINDOWS,
     reason="windows prevents two connections to the same file even in the same process",
 )
+
+
+@pytest.fixture(scope="session")
+def data_dir() -> Path:
+    """Return the test data directory.
+
+    Returns
+    -------
+    Path
+        Test data directory
+    """
+    root = Path(__file__).absolute().parents[1]
+
+    return root / "ci" / "ibis-testing-data"

ibis/tests/benchmarks/test_benchmarks.py

@@ -1019,3 +1019,34 @@ def test_dedup_schema(benchmark):
             itertools.cycle(("int", "string", "array<int>", "float")),
         ),
     )
+
+
+@pytest.fixture(scope="session")
+def pgtable(data_dir):
+    pd = pytest.importorskip("pandas")
+    pytest.importorskip("psycopg")
+
+    from ibis.backends.postgres.tests.conftest import (
+        IBIS_TEST_POSTGRES_DB,
+        PG_HOST,
+        PG_PASS,
+        PG_PORT,
+        PG_USER,
+    )
+
+    con = ibis.postgres.connect(
+        user=PG_USER,
+        password=PG_PASS,
+        host=PG_HOST,
+        port=PG_PORT,
+        database=IBIS_TEST_POSTGRES_DB,
+    )
+    name = ibis.util.gen_name("functional_alltypes_bench")
+    yield con.create_table(
+        name, obj=pd.read_csv(data_dir / "csv" / "functional_alltypes.csv"), temp=True
+    )
+    con.disconnect()
+
+
+def test_postgres_record_batches(pgtable, benchmark):
+    benchmark(pgtable.to_pyarrow)