You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Share the buffer refill between fetchone() and the bulk fetches of the result sets #1193
AthenaResultSet, AthenaArrowResultSet, and AthenaS3FSResultSet buffer converted rows in self._rows and refill it with _fetch().
They differ only in when fetchone() refills the buffer:
The base fetchone() (pyathena/result_set.py) calls _fetch() only while NextToken remains, because the base _fetch() raises ProgrammingError without one.
AthenaArrowResultSet.fetchone() and AthenaS3FSResultSet.fetchone() override it only to call _fetch() whenever the buffer is empty; their _fetch() returns without rows at the end of the results.
Since #1187, the base fetchmany() and fetchall() (_take_rows()) and the asyncio _async_take_rows() take buffered rows in bulk and call fetchone() once whenever the buffer is empty, only to trigger that refill.
As a result:
The first row of each page or batch takes a separate path through fetchone() in the bulk take.
Proposed change
Move the refill condition into one place, so that fetchone() and the bulk take share it, and remove the Arrow and S3FS fetchone() overrides:
Make the base _fetch() return without rows when no NextToken remains, instead of raising ProgrammingError, so that every result set's _fetch() means "read the next rows, if any". fetchone() then calls _fetch() whenever the buffer is empty.
Alternatively, keep _fetch() as is and add a hook (for example _fetch_next()): if self._next_token: self._fetch() in the base, self._fetch() in Arrow and S3FS.
Either way, _take_rows() could refill the buffer directly and take every row in bulk.
A subclass whose fetchone() does not use the buffer would then need to override _take_rows() as well, as the pandas and Polars result sets already do; the PyAthena result sets have no such fetchone().
This changes the roles of private methods in the result set classes, so the approach needs agreement first.
No public behavior changes: the rows, rownumber, and the GetQueryResults requests stay the same.
Validation plan (if implementing)
Unit tests in tests/pyathena/test_result_set.py and tests/pyathena/aio/test_result_set.py for fetches across pages, after close(), and on a failed page request, plus tests of the Arrow and S3FS refill at the end of the results.
The fetch tests of the default, Arrow, and S3FS cursors and their asyncio counterparts against Athena.
Activity