Repository: ibis-project/ibis (Python)

Issue:
feat: Improve `to_pyarrow_batches` for PostgreSQL backend
### Is your feature request related to a problem?

Hi,

It seems that `to_pyarrow_batches` is implemented somewhat naively in many backends. In many cases (including the SQL backends) all the data is first instantiated in the client-side cursor (or as a pandas DF) and then partitioned to batches. This means that something like

`remote_con.table('huge_table').to_pyarrow_batches(...)`  tries to allocate the whole table in memory.

### What is the motivation behind your request?

_No response_

### Describe the solution you'd like

PostgreSQL (and maybe other backends) has a mechanism to batch the results server-side.

I can make a PR for this.

### What version of ibis are you running?

10.2.0

### What backend(s) are you using, if any?

PostgreSQL

### Code of Conduct

- [x] I agree to follow this project's Code of Conduct


Discussion: @ronif Thanks for the issue!

Would definitely review a PR to improve `to_pyarrow_batches()` for Postgres.



Fix (patch):
diff --git a/ibis/backends/postgres/__init__.py b/ibis/backends/postgres/__init__.py
index 3f38dd718db2..c41125ea8a3e 100644
--- a/ibis/backends/postgres/__init__.py
+++ b/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))