Skip to content

Commit

Permalink
update read_database now frame init can take alchemy rows "as-is"
Browse files Browse the repository at this point in the history
  • Loading branch information
alexander-beedie committed Nov 15, 2024
1 parent 2c57ea7 commit 7657872
Showing 1 changed file with 13 additions and 9 deletions.
22 changes: 13 additions & 9 deletions py-polars/polars/io/database/_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -178,24 +178,24 @@ def _fetch_arrow(
yield arrow

@staticmethod
def _fetchall_rows(result: Cursor) -> Iterable[Sequence[Any]]:
def _fetchall_rows(result: Cursor, *, is_alchemy: bool) -> Iterable[Sequence[Any]]:
"""Fetch row data in a single call, returning the complete result set."""
rows = result.fetchall()
return (
[tuple(row) for row in rows]
if rows and not isinstance(rows[0], (list, tuple, dict))
else rows
rows
if rows and (is_alchemy or isinstance(rows[0], (list, tuple, dict)))
else [tuple(row) for row in rows]
)

def _fetchmany_rows(
self, result: Cursor, batch_size: int | None
self, result: Cursor, *, batch_size: int | None, is_alchemy: bool
) -> Iterable[Sequence[Any]]:
"""Fetch row data incrementally, yielding over the complete result set."""
while True:
rows = result.fetchmany(batch_size)
if not rows:
break
elif isinstance(rows[0], (list, tuple, dict)):
elif is_alchemy or isinstance(rows[0], (list, tuple, dict)):
yield rows
else:
yield [tuple(row) for row in rows]
Expand Down Expand Up @@ -267,7 +267,7 @@ def _from_rows(
self.result = _run_async(self.result)
try:
if hasattr(self.result, "fetchall"):
if self.driver_name == "sqlalchemy":
if is_alchemy := (self.driver_name == "sqlalchemy"):
if hasattr(self.result, "cursor"):
cursor_desc = [
(d[0], d[1:]) for d in self.result.cursor.description
Expand Down Expand Up @@ -297,9 +297,13 @@ def _from_rows(
orient="row",
)
for rows in (
self._fetchmany_rows(self.result, batch_size)
self._fetchmany_rows(
self.result,
batch_size=batch_size,
is_alchemy=is_alchemy,
)
if iter_batches
else [self._fetchall_rows(self.result)] # type: ignore[list-item]
else [self._fetchall_rows(self.result, is_alchemy=is_alchemy)] # type: ignore[list-item]
)
)
return frames if iter_batches else next(frames) # type: ignore[arg-type]
Expand Down

0 comments on commit 7657872

Please sign in to comment.