From 89ca378b1f71c444a10cde47e81baa88ffe25b80 Mon Sep 17 00:00:00 2001 From: Alexander Beedie Date: Wed, 16 Oct 2024 10:10:38 +0400 Subject: [PATCH 1/2] fix(python): Ensure that `read_database` takes advantage of Arrow return from a `duckdb_engine` connection when using a SQLAlchemy `Selectable` --- py-polars/polars/io/database/_executor.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/py-polars/polars/io/database/_executor.py b/py-polars/polars/io/database/_executor.py index 278e3e8e0738..7111733ced1c 100644 --- a/py-polars/polars/io/database/_executor.py +++ b/py-polars/polars/io/database/_executor.py @@ -392,7 +392,7 @@ def _normalise_cursor(self, conn: Any) -> Cursor: return conn.engine.raw_connection().cursor() elif conn.engine.driver == "duckdb_engine": self.driver_name = "duckdb" - return conn.engine.raw_connection().driver_connection + return conn elif self._is_alchemy_engine(conn): # note: if we create it, we can close it self.can_close_cursor = True @@ -505,8 +505,11 @@ def execute( ) result = cursor_execute(query, *positional_options) - # note: some cursors execute in-place + # note: some cursors execute in-place, some access results via an additional property result = self.cursor if result is None else result + if self.driver_name == "duckdb": + result = result.cursor + self.result = result return self From 2e207ae3ee9cf709d45cf3a715521bc97872bd66 Mon Sep 17 00:00:00 2001 From: Alexander Beedie Date: Wed, 16 Oct 2024 10:45:11 +0400 Subject: [PATCH 2/2] fix line-length (lint) --- py-polars/polars/io/database/_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/py-polars/polars/io/database/_executor.py b/py-polars/polars/io/database/_executor.py index 7111733ced1c..3162c974a317 100644 --- a/py-polars/polars/io/database/_executor.py +++ b/py-polars/polars/io/database/_executor.py @@ -505,7 +505,7 @@ def execute( ) result = cursor_execute(query, *positional_options) - # note: some cursors execute in-place, some access results via an additional property + # note: some cursors execute in-place, some access results via a property result = self.cursor if result is None else result if self.driver_name == "duckdb": result = result.cursor