Skip to content

PolarsCursor with chunksize returns 0 rows instead of raising when the CSV read fails #820

Description

@laughingman7743

Problem

With PolarsCursor and chunksize set, a failure while reading the CSV query result is not raised to the caller.
The chunk iterator ends normally, and the cursor returns 0 rows for a query whose result contains 10,000,000 rows.
Callers cannot distinguish this from an empty result, so data is lost silently.

Expected: the read failure is raised (for example as OperationalError, as the non-chunked path does), and no partial or empty result is returned as success.

Reproduction

Observed during the #644 large-data benchmark on master 0c26c7e:

  • Query result: CSV of 10,000,000 rows (2,846,664,276 bytes), read with PolarsCursor (sync, thread, and aio APIs), chunksize=100000.
  • The Polars download failed with polars.exceptions.ComputeError: failed to reserve 2846664276 bytes on disk to download uri = s3://.../<query-id>.csv: Os { code: 28, kind: StorageFull, ... } (the disk-cache growth is tracked in PolarsCursor chunked CSV reads download the whole result into Polars' local file cache #821).
  • The query had SUCCEEDED, and the CSV object in S3 was complete (2,846,664,276 bytes).
  • The benchmark consumed the cursor to exhaustion and got 0 rows in 0.3 seconds, with no exception. The error appeared only on stderr as:
Exception ignored in: <function LazyFrame.collect_batches.<locals>.BatchCollector.__del__ at 0x...>
Traceback (most recent call last):
  File ".../polars/lazyframe/frame.py", line 3889, in __del__
    self._fut.result()
  ...
polars.exceptions.ComputeError: failed to reserve 2846664276 bytes on disk to download uri = ...

Cause, from the Polars 1.36.1 source of LazyFrame.collect_batches(): the background task runs sink_batches and, in a finally block, always puts the end-of-batches sentinel None into the queue.
BatchCollector.__next__ turns that sentinel into StopIteration, so the iterator ends normally even when sink_batches raised.
The exception is only re-raised by self._fut.result() in BatchCollector.__del__, where Python reports it as "Exception ignored".

AthenaPolarsResultSet._iter_csv_chunks (pyathena/polars/result_set.py:663) iterates lazy_df.collect_batches(chunk_size=self._chunksize) inside try/except Exception, so the exception never reaches that handler.
_iter_parquet_chunks (pyathena/polars/result_set.py:694) uses the same collect_batches iteration for UNLOAD results and is expected to have the same exposure; that path has not been reproduced yet.

Environment

  • PyAthena master 0c26c7e (3.36.1.dev240), Python 3.12.14, Amazon Linux 2023 on EC2 r7i.2xlarge
  • polars 1.36.1 (from the repository lockfile); later Polars versions have not been checked
  • Cursor: PolarsCursor and AioPolarsCursor, chunksize set, CSV results

Proposed fix (optional)

Do not rely on collect_batches() iteration alone to report failures. Candidates to discuss:

  • Drive LazyFrame.sink_batches from PyAthena's own producer thread and queue, and re-raise the producer's exception from the consumer.
  • Or, after the iterator is exhausted, surface the background failure explicitly before reporting completion.

A regression test can use a scan that fails partway (for example, an unreadable object) and assert that iterating the chunked result raises instead of yielding zero rows.

Found while measuring #644.

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

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions