diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 508f644d7..cd471a103 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -197,24 +197,27 @@ def converters(self) -> dict[str, Callable[[str | None], Any | None]]: @override def _fetch(self) -> None: - try: - rows = next(self._batches) - except StopIteration: + """Buffer the rows of the next record batch that has any. + + A table can hold empty record batches between others, such as one of + an empty result file, which would otherwise end the fetches early. + """ + rows = next((batch for batch in self._batches if batch.num_rows), None) + if rows is None: return + # Read the columns and their converters by position; to_pydict() and the + # converters property keep one column per name. + columns = [column.to_pylist() for column in rows.columns] + description = self.description if self.description else [] + converters = [self._converter.get(d[1]) for d in description] + if any(convert is not _to_default for convert in converters): + processed_rows = [ + tuple(convert(v) for convert, v in zip(converters, row, strict=False)) + for row in zip(*columns, strict=False) + ] else: - # Read the columns and their converters by position; to_pydict() and the - # converters property keep one column per name. - columns = [column.to_pylist() for column in rows.columns] - description = self.description if self.description else [] - converters = [self._converter.get(d[1]) for d in description] - if any(convert is not _to_default for convert in converters): - processed_rows = [ - tuple(convert(v) for convert, v in zip(converters, row, strict=False)) - for row in zip(*columns, strict=False) - ] - else: - processed_rows = list(zip(*columns, strict=False)) - self._rows.extend(processed_rows) + processed_rows = list(zip(*columns, strict=False)) + self._rows.extend(processed_rows) def _open_csv(self, location: str) -> Any: """Open the CSV result file for ``pyarrow.csv.read_csv()``. diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 2762b9b86..8592c7aa1 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -5,6 +5,7 @@ import csv import importlib import logging +import sys import threading from collections import abc from collections.abc import Callable, Iterable, Iterator @@ -992,8 +993,9 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] try: - # extend() keeps the rows that it appended before an exception. - rows.extend(islice(self._row_iter, size)) + # extend() keeps the rows that it appended before an exception, and + # islice() takes no size above sys.maxsize. + rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize))) finally: if rows: self._count_rows(len(rows)) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 74a2e099f..b79802547 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -12,6 +12,7 @@ import csv import importlib.util import logging +import sys import threading from collections import abc from collections.abc import Callable, Iterable, Iterator @@ -647,8 +648,9 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] try: - # extend() keeps the rows that it appended before an exception. - rows.extend(islice(self._row_iter, size)) + # extend() keeps the rows that it appended before an exception, and + # islice() takes no size above sys.maxsize. + rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize))) finally: if rows: self._count_rows(len(rows)) diff --git a/pyathena/result_set.py b/pyathena/result_set.py index c0a7cf9e3..1e1445641 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -5,7 +5,6 @@ import collections import logging import operator -import sys from datetime import datetime from itertools import islice from typing import ( @@ -529,8 +528,7 @@ def fetchmany( """ if not size or size <= 0: size = self._arraysize - # islice() in the subclasses' _take_rows() takes no size above sys.maxsize. - return self._take_rows(min(operator.index(size), sys.maxsize)) + return self._take_rows(operator.index(size)) @override def fetchall( diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index 46be7dc41..cd214b6b4 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -64,6 +64,53 @@ def test_fetch_across_batches(self): assert result_set.fetchone() is None assert result_set.rownumber == 3 + @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall", "iter"]) + @pytest.mark.parametrize( + ("chunks", "expected"), + [ + ([[1, 2, 3], [], [4]], [1, 2, 3, 4]), + ([[], [1, 2]], [1, 2]), + ([[1], [], [], [2]], [1, 2]), + ], + ids=["middle", "leading", "consecutive"], + ) + def test_fetch_across_empty_batches(self, method, chunks, expected): + """Empty record batches before or between others do not end the fetches. + + A table read from result files that include an empty one has such + batches; the fetches used to stop at the first one. to_batches() yields + no empty batch after the last rows. No AWS calls; the first page + request and the result table are mocked. + """ + table = pa.concat_tables([pa.table({"a": pa.array(chunk, pa.int64())}) for chunk in chunks]) + with ( + patch.object(AthenaArrowResultSet, "_pre_fetch"), + patch.object(AthenaArrowResultSet, "_create_s3_file_system"), + patch.object(AthenaArrowResultSet, "_as_arrow", return_value=table), + ): + result_set = AthenaArrowResultSet( + connection=MagicMock(), + converter=DefaultArrowTypeConverter(), + query_execution=MagicMock( + state=AthenaQueryExecution.STATE_SUCCEEDED, + output_location="s3://bucket/query.csv", + ), + arraysize=2, + retry_config=RetryConfig(), + ) + assert 0 in [batch.num_rows for batch in table.to_batches(2)] + if method == "fetchone": + rows = list(iter(result_set.fetchone, None)) + elif method == "fetchmany": + rows = result_set.fetchmany(10) + elif method == "fetchall": + rows = result_set.fetchall() + else: + rows = list(result_set) + assert rows == [(value,) for value in expected] + assert result_set.rownumber == len(expected) + assert result_set.fetchone() is None + def test_read_csv_timestamps_with_the_same_name(self): """Timestamp columns are converted by position, also with names that repeat. diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index 83df71187..ec647d1bf 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -20,13 +20,6 @@ from tests.pyathena.util import close_during_next_request, get_query_results_responses -def _page(values, next_token=None): - response = {"ResultSet": {"Rows": [{"Data": [{"VarCharValue": v}]} for v in values]}} - if next_token: - response["NextToken"] = next_token - return response - - def _row(*values): return {"Data": [{} if v is None else {"VarCharValue": v} for v in values]} @@ -180,7 +173,7 @@ def test_fetch_all_rows_skips_column_labels_only_on_first_page(self): result_set._process_metadata( {"ResultSet": {"ResultSetMetadata": {"ColumnInfo": [{"Name": "a", "Type": "varchar"}]}}} ) - pages = [_page(["a", "1"], "token"), _page(["a", "2"])] + pages = get_query_results_responses(["a", "1"], ["a", "2"]) with patch.object(result_set, "_get_query_results", side_effect=pages): assert result_set._fetch_all_rows() == [("1",), ("a",), ("2",)]