Skip to content

PolarsDataFrameIterator.close() does not stop a whole-DataFrame iterator #1019

Description

@laughingman7743

Problem

PolarsDataFrameIterator.close() only closes a generator reader (pyathena/polars/result_set.py, close() checks isinstance(self._reader, GeneratorType)).
When the iterator wraps a whole DataFrame (iter([df])), which happens without chunksize, with managed query result storage, and for failed or empty results, close() does nothing, so the iterator still yields the DataFrame afterwards and keeps the reference to it.

import polars as pl
from pyathena.polars.result_set import PolarsDataFrameIterator

chunks = PolarsDataFrameIterator(pl.DataFrame({"a": [1]}), {}, ["a"])
chunks.close()
next(chunks)  # still returns the DataFrame

PandasDataFrameIterator.close() replaces its reader with iter(()) for every reader kind, so a closed pandas iterator is exhausted.

Expected

After close(), a PolarsDataFrameIterator yields nothing and releases its reader for every reader kind, as PandasDataFrameIterator does.

Found by the independent review of #1005 (#935). Low impact: AthenaPolarsResultSet.close() replaces its own iterator anyway; this affects iterators that callers keep from iter_chunks().

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions