From ffe076fc8d8cc437c96499f0038ca2c499a54ac5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Fri, 5 Jan 2024 17:13:45 +0545 Subject: [PATCH 1/3] ENH: enable server side cursors when chunksize is set. Closes #35689. --- pandas/io/sql.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pandas/io/sql.py b/pandas/io/sql.py index 3a58daf681cfb..1763400198240 100644 --- a/pandas/io/sql.py +++ b/pandas/io/sql.py @@ -1658,6 +1658,8 @@ def run_transaction(self): def execute(self, sql: str | Select | TextClause, params=None): """Simple passthrough to SQLAlchemy connectable""" args = [] if params is None else [params] + if self.returns_generator: + self.con.execution_options(stream_results=True) if isinstance(sql, str): return self.con.exec_driver_sql(sql, *args) return self.con.execute(sql, *args) @@ -1836,11 +1838,11 @@ def read_query( read_sql """ + self.returns_generator = chunksize is not None result = self.execute(sql, params) columns = result.keys() if chunksize is not None: - self.returns_generator = True return self._query_iterator( result, self.exit_stack, From 07b363a2d86c056924c6fae731fe4682fb39e70c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Thu, 11 Jan 2024 23:18:00 +0545 Subject: [PATCH 2/3] add simple test for execution_options --- pandas/tests/io/test_sql.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/pandas/tests/io/test_sql.py b/pandas/tests/io/test_sql.py index 6645aefd4f0a7..33d4b100146d6 100644 --- a/pandas/tests/io/test_sql.py +++ b/pandas/tests/io/test_sql.py @@ -1189,6 +1189,22 @@ def test_read_iris_table_chunksize(conn, request): check_iris_frame(iris_frame) +@pytest.mark.parametrize("conn", sqlalchemy_connectable_iris) +def test_sqlalchemy_read_table_chunksize_stream_results(conn, request): + conn = request.getfixturevalue(conn) + with sql.SQLDatabase(conn) as pandasSQL: + assert list(pandasSQL.read_table("iris", chunksize=10000)) + assert pandasSQL.con.get_execution_options() == {"stream_results": True} + + +@pytest.mark.parametrize("conn", sqlalchemy_connectable_iris) +def test_sqlalchemy_read_sql_chunksize_stream_results(conn, request): + conn = request.getfixturevalue(conn) + with sql.SQLDatabase(conn) as pandasSQL: + assert list(pandasSQL.read_sql("SELECT * FROM iris", chunksize=10000)) + assert pandasSQL.con.get_execution_options() == {"stream_results": True} + + @pytest.mark.parametrize("conn", sqlalchemy_connectable) def test_to_sql_callable(conn, test_frame1, request): conn = request.getfixturevalue(conn) From ed5611b62f0e7be93313b3c5956a188bb65ee739 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Fri, 12 Jan 2024 10:56:39 +0545 Subject: [PATCH 3/3] fix tests --- pandas/tests/io/test_sql.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pandas/tests/io/test_sql.py b/pandas/tests/io/test_sql.py index 33d4b100146d6..bfd9fa06c8438 100644 --- a/pandas/tests/io/test_sql.py +++ b/pandas/tests/io/test_sql.py @@ -1194,7 +1194,8 @@ def test_sqlalchemy_read_table_chunksize_stream_results(conn, request): conn = request.getfixturevalue(conn) with sql.SQLDatabase(conn) as pandasSQL: assert list(pandasSQL.read_table("iris", chunksize=10000)) - assert pandasSQL.con.get_execution_options() == {"stream_results": True} + execution_options = pandasSQL.con.get_execution_options() + assert execution_options["stream_results"] is True @pytest.mark.parametrize("conn", sqlalchemy_connectable_iris) @@ -1202,7 +1203,8 @@ def test_sqlalchemy_read_sql_chunksize_stream_results(conn, request): conn = request.getfixturevalue(conn) with sql.SQLDatabase(conn) as pandasSQL: assert list(pandasSQL.read_sql("SELECT * FROM iris", chunksize=10000)) - assert pandasSQL.con.get_execution_options() == {"stream_results": True} + execution_options = pandasSQL.con.get_execution_options() + assert execution_options["stream_results"] is True @pytest.mark.parametrize("conn", sqlalchemy_connectable)