Skip to content

Share the buffer refill between fetchone() and the bulk fetches of the result sets #1193

Closed

Description

@laughingman7743

Use case

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 same fetchone() body exists in the base, Arrow, and S3FS result sets (and in AthenaAioResultSet with _async_fetch()), and Take buffered rows in bulk in fetchmany() and fetchall() #1187 had to change each copy the same way.
  • 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:

  1. 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.
  2. 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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions