Skip empty Arrow record batches and move the fetchmany() clamp

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()``.

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))

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))

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(

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.
 

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",)]