Keep the whole pandas and Polars results reusable

[laughingman7743]

WHAT

Without chunksize, the pandas and Polars result sets keep the whole result as a DataFrame and no longer hand out the iterator that the fetch methods read.

• PandasCursor.as_pandas(), PolarsCursor.as_polars(), and PolarsCursor.as_arrow() (and the async/aio variants, which share the result sets) return the whole result on every call. as_pandas() / as_polars() return the same DataFrame object each time; as_arrow() converts it.
• The fetch methods read a shallow copy of that DataFrame (copy(deep=False) under pandas 3 Copy-on-Write, clone() for Polars), so assignments to the returned DataFrame do not change the fetched rows, and calling as_*() does not move the fetch position. Mutable values inside cells, such as lists from JSON columns, stay shared; a deep copy would double the memory.
• iter_chunks() without chunking returns a new iterator over the stored DataFrame on each call, so it no longer drains the fetch iterator, and PandasCursor.iter_chunks() closing its iterator no longer empties the result.
• Chunked reads (explicit chunksize, or a chunk size chosen by auto_optimize_chunksize) keep their single-pass iterator shared with the fetch methods; the docstrings now state this.
• A result that is not read in chunks despite chunksize (pandas UNLOAD results, which read_parquet reads whole, and results without an S3 result file) also gets the stored DataFrame; for pandas, as_pandas() with chunksize then returns a new iterator over it on each call.
• The date/time truncation for pandas is applied once when the result is read, instead of when the single chunk is first yielded.
• GetQueryResults fallback (no S3 result file, e.g. managed query result storage): the rows are already converted, so PandasCursor no longer applies the time truncation to them (a TIME column raised AttributeError on fetch/as_pandas()), and PolarsCursor no longer converts them again in the fetch methods (a TIME column raised TypeError: strptime() argument 1 must be str, not datetime.time). Found by the independent review; the pandas case would otherwise have moved into execute() with this change.
• AthenaPolarsResultSet.close() closes the chunk iterator before replacing it, as the pandas result set does, so a retained iter_chunks() iterator stops.

Release notes (4.0.0):

• Behavior change: without chunksize, a repeated as_pandas() / as_polars() / as_arrow() returns the whole result instead of an empty frame, and fetches after them return the rows; as_pandas() / as_polars() return the same object on each call, and the result set keeps it until close() or the next execute().
• Fix: with managed query result storage, TIME columns no longer fail in PandasCursor and in PolarsCursor fetches.
• Fix: closing a PolarsCursor result set closes its chunk iterator.

WHY

Closes #935.

as_pandas() returned next() / joined the iterator that fetchone() / fetchmany() / fetchall() also read, so a second call returned an empty DataFrame and a fetch afterwards returned nothing (pandas since v2.13.0; Polars since v3.24.0, which replaced the stored _df with the iterator wrapper).

Maintainer decisions: the auto_optimize_chunksize chunked case keeps the single-pass behavior from #950 (no caching, no S3 re-read), the non-chunked iter_chunks() is aligned in this PR, and the GetQueryResults fallback conversion and Polars close() defects found in review are fixed here. The chunked Polars UNLOAD metadata defect found in the same review is #1007.

TEST

Tested commit: c857dd3 (7ce2bf4 changes one comment only; just lint passed)

• just lint: passed.
• New integration tests TestPandasCursor::test_pandas_cursor_whole_result_reused and TestPolarsCursor::test_whole_result_reused cover repeated as_*(), assignment isolation, fetch position across as_*() / iter_chunks(), and repeated iter_chunks(). Both fail on the original code (as_pandas() / as_polars() returned an empty frame on the second call).
• test_fetch_all_rows[managed] (pandas and Polars) now selects a TIME column, and the new TestPolarsCursor::test_close_stops_chunks checks that a retained chunk iterator stops after close(). With the source changes of c857dd3 reverted, these three fail (AttributeError, TypeError, and remaining chunks); with them, they pass.
• uv run --env-file .env pytest -n 4 tests/pyathena/pandas tests/pyathena/polars tests/pyathena/aio/pandas tests/pyathena/aio/polars -k "as_pandas or as_polars or as_arrow or iter_chunks or fetch or close or managed or Iterator or whole_result": 177 passed (sync, async, and aio cursors, including the managed work group).
• Offline check (pandas 3.0.6, polars 1.44.2): without the shallow copy / clone, df["a"] = 0, an in-place rename, or df[0, "a"] = 0 on the returned DataFrame changed the rows that the fetch iterator later returned; with them, it did not.
• Not run locally: the full just test pyathena suite; it runs in CI when the PR is marked Ready.

🤖 Generated with Claude Code

[laughingman7743]

Self-review round one (implementation behavior): CLEAN

Base a18ebdaa80c0c85466eacbdb3eb3df76d7ed135a, head f41ca3c05cf61338ab2ab408541efdb6cf13b9fc.

Covered: AthenaPandasResultSet / AthenaPolarsResultSet construction for the S3 output, GetQueryResults fallback, and failed-state paths; as_pandas(), as_polars(), as_arrow(), iter_chunks(), close(); the sync/async/aio cursor wrappers (no subclasses; PandasCursor.iter_chunks() closes the iterator it receives, which is now a fresh one in non-chunked mode); explicit chunksize and auto_optimize_chunksize chunked paths (unchanged single-pass); the two new tests.

Checks:

• Hypothesis: eager _trunc_date in __init__ could now raise on results that were never read. Refuted live: zero-row time, all-NULL time, and mixed time/NULL results convert identically on master and this head (master only differs in returning [] from fetchall() after as_pandas(), which is the bug).
• Mutation isolation verified offline (pandas 3.0.6 CoW, polars 1.44.2): without the shallow copy / clone, df["a"] = 0, in-place rename, and df[0, "a"] = 0 leaked into fetched rows; with them, they do not.
• Both new tests fail on the original code.

Limitation (not a defect): in non-chunked mode the result set now holds the DataFrame until close() or the next execute(), as before v2.13.0 for pandas, before v3.24.0 for Polars, and like ArrowCursor. Previously a fetch-only caller released it once the single chunk was exhausted.

[laughingman7743]

Self-review round two (claims, callers, operations): FINDINGS → repaired in a26c763

Base a18ebdaa80c0c85466eacbdb3eb3df76d7ed135a, head f41ca3c05cf61338ab2ab408541efdb6cf13b9fc. Full pass over the PR body, the commit message, and the changed docstrings and comments.

Claims checked:

• "pandas since v2.13.0; Polars since v3.24.0": git tag --contains puts d5cb976 first in v2.13.0 and 6f7addc first in v3.24.0. Before v3.24.0, as_polars() returned self._df.
• Shallow-copy isolation relies on Copy-on-Write: pyproject.toml requires pandas>=3.0.0 (CoW is always on) and polars>=1.39.0; verified offline in round one.
• "async/aio variants share the result sets": pandas/async_cursor.py, polars/async_cursor.py, and aio/{pandas,polars}/cursor.py construct AthenaPandasResultSet / AthenaPolarsResultSet.
• PandasCursor.iter_chunks() closes its iterator (with result_set.iter_chunks() as chunks); on master that closed the shared fetch iterator.

Findings (repaired):

1. pyathena/pandas/result_set.py iter_chunks(): the new sentence said chunks with chunksize come from the fetch iterator, but _read_parquet ignores chunksize. A pandas UNLOAD result with chunksize=1000 gets the stored DataFrame and a fresh iterator. The docstring now conditions on the CSV result being read in chunks.
2. pyathena/polars/result_set.py as_polars() / as_arrow() / iter_chunks(): "when chunksize is set ... a later call returns an empty DataFrame" was false for results without an S3 result file (the GetQueryResults path stores the DataFrame even with chunksize). The wording now conditions on the result file being read in chunks.
3. The PR body now states the pandas UNLOAD/chunksize case.

Callers and operations: no change in S3 or Athena requests (same reads, same places). No user guide in docs/ describes the old consuming behavior. #927 wrote the guides for the intended behavior. Behavior changes are listed in the release note. Limitation from round one (DataFrame retained until close()) still stands.

[laughingman7743]

Independent review (relayed): FINDINGS

Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort high, sandbox read-only, session 01a1015d-691c-71d1-8c1a-e2e76abd7851. Static review only (no builds, tests, or network). Snapshot: detached worktree at head a26c7633d81164fed18507525713ecb950385fc4, base a18ebdaa80c0c85466eacbdb3eb3df76d7ed135a. The prompt contained the literal diff and the intended behavior, without the PR number, description, commits, or prior findings. Afterwards, both the snapshot and the PR worktree were clean and still at a26c7633.

Covered (reviewer): S3 CSV, UNLOAD Parquet, API fallback, empty and non-SUCCEEDED construction; whole-result access, explicit/automatic chunking, fetch positions, close/reset, CSV streams and readers; mutation isolation, memory retention, pandas time truncation, docs; sync, future-based async, and aio wrappers; both new tests.

Findings and author verification:

1. Diff-related, P2 pyathena/pandas/result_set.py:364: the shallow copy does not isolate mutable cell contents. With CAST(ARRAY[1, 2] AS JSON), df.at[0, "j"].append(3) changes the later fetchone(). Verified (object cells are shared Python objects). Action: narrow the comment and the PR body to assignments to the DataFrame. A deep copy would double the memory.
2. Pre-existing, made worse by the diff, P2 pyathena/pandas/result_set.py:358: in the GetQueryResults fallback, TIME values are already datetime.time, so _trunc_date's .dt.time raises AttributeError. Verified live with the managed work group: master raises it on fetchone(), while this head raises it in execute(). Action: fix in this PR.
3. Pre-existing, P2 pyathena/polars/result_set.py:138: the fallback DataFrame already holds converted values, and iterrows() converts them again. Verified live: TIME fetchone() raises TypeError: strptime() argument 1 must be str, not datetime.time on both master and head. Action (maintainer chose to fold in): fix in this PR.
4. Pre-existing, P2 pyathena/polars/result_set.py:359: chunked UNLOAD keeps the UNLOAD command metadata. Verified live: unload=True, chunksize=10 gives description [('rows', 'bigint', ...)] and fetchone() (None,); without chunksize, (1, 'x'). Action (maintainer): separate issue.
5. Pre-existing, P2 pyathena/polars/result_set.py:711: close() replaces _df_iter without closing it (pandas closes it). Verified by reading the code. Action (maintainer chose to fold in): fix in this PR.

[laughingman7743]

Repair: c857dd3 (on top of a26c763; same base a18ebda)

• Finding 1: narrowed both comments ("assignments to the DataFrame"; pandas notes that mutable cell values such as JSON lists stay shared) and the PR body. No deep copy.
• Finding 2: the pandas GetQueryResults fallback no longer applies _trunc_date. DefaultTypeConverter already returns datetime.time for time. time with time zone stays a string, as in Cursor; the old .dt.time also raised on it.
• Finding 3: the Polars fallback uses no row converters (_df_converters = {}), for the fetch iterator and for iter_chunks().
• Finding 4: filed as PolarsCursor with unload=True and chunksize returns the UNLOAD statement's metadata and None rows #1007.
• Finding 5: AthenaPolarsResultSet.close() calls self._df_iter.close() before replacing it. For a whole-result iterator, close() is a no-op; a chunk generator gets GeneratorExit, which the except Exception in _iter_*_chunks does not catch.

Validation: just lint passed. test_fetch_all_rows[managed] (pandas and Polars) now selects a TIME column, and the new test_close_stops_chunks covers finding 5. With the source changes of c857dd3 reverted, these 3 fail; with them, they pass. The related pandas, Polars, and aio subsets: 177 passed.

Self-review of the repair (round one, behavior): no findings. The fallback DataFrame types are unchanged, and only the fetch conversion and the truncation differ. Round two (claims): the commit message, comments, and PR body were rechecked against the measured errors (AttributeError in execute() on a26c763 and in fetchone() on master; TypeError from strptime()); no findings.

[laughingman7743]

Independent follow-up review (relayed): no regressions; 2 pre-existing findings

Reviewer: Codex CLI 0.160.0, gpt-6-astra, effort high, read-only, session 01a1016a-31a5-7862-ac29-da0c5dd01643. Static review of a26c7633..c857dd3e06a2d5a466e2997f110ea533a2267db5 (range-diff: earlier commits unchanged) traced into both result sets, the cursor wrappers, _fetch_all_rows, and the converters. Afterwards, the snapshot and the PR worktree were clean at c857dd3e.

The reviewer resolved all four earlier comments: assignment isolation is documented with the shared mutable cells; fallback TIME bypasses .dt.time (also with chunksize); Polars fallback rows are not converted again, for fetch and iter_chunks().iterrows(); and the stored CSV/Parquet chunk generators are closed. By source tracing, the managed TIME tests and test_close_stops_chunks fail without the fix.

Pre-existing findings and author disposition:

1. P2 pyathena/converter.py:481: DefaultTypeConverter has no time with time zone mapping, so fallback values stay strings. That is the existing Cursor behavior, out of scope. It did show that the new comment at pyathena/pandas/result_set.py:358 ("time columns hold times") was too broad. Repaired in 7ce2bf4 (comment only: "already converted and need no time truncation"). just lint passed.
2. P3 pyathena/polars/result_set.py:124: PolarsDataFrameIterator.close() does nothing for a whole-DataFrame reader, unlike the pandas iterator. Deferred: this is pre-existing, the iterator only holds a DataFrame reference, and AthenaPolarsResultSet.close() replaces the iterator anyway.

[laughingman7743]

Independent follow-up (relayed): CLEAN for c857dd3e..7ce2bf4ca262003d4b3708a77161a7c03c0382e6 (comment only).

Reviewer: Codex CLI 0.160.0, gpt-6-astra, effort high, read-only, session 01a1016e-7848-7901-ac2a-c2eb4a17074f. Result: the comment at pyathena/pandas/result_set.py:358 is accurate. With no output location, DefaultTypeConverter turns time into datetime.time and leaves time with time zone as a string; neither needs _trunc_date(). Static review only.