From 20ead4675b6de12742a4655ec3010799048690b3 Mon Sep 17 00:00:00 2001 From: Alexander Beedie Date: Wed, 16 Oct 2024 11:49:27 +0400 Subject: [PATCH] fix(python): Ensure that `read_database` takes advantage of Arrow return from a `duckdb_engine` connection when using a SQLAlchemy `Selectable` (#19255) --- 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..3162c974a317 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 a property result = self.cursor if result is None else result + if self.driver_name == "duckdb": + result = result.cursor + self.result = result return self