Skip to content

Pre-existing UNLOAD read defects found in the #1214 review #1235

Closed

Description

@laughingman7743

Found while reviewing #1233 (#1214). These defects predate it, are still present on master d84526e, and #1233 does not change them.

1. PolarsCursor with use_pyarrow=True cannot read UNLOAD results

_read_parquet() passes the unload location, an S3 prefix, to pl.read_parquet(). With use_pyarrow=True, Polars opens a remote string as one file through fsspec.open(), so the read fails before decoding.

Measured live on master 9e2bf6e with Polars 2.0.0:

cursor = connect(cursor_class=PolarsCursor).cursor(unload=True)
cursor.execute("SELECT 1 AS id", use_pyarrow=True)
# OperationalError: <bucket>/<staging>/unload/<date>/<uuid>
# caused by FileNotFoundError on the prefix

docs/polars.md says that the other reads use Polars' native S3 access by default, and fsspec with use_pyarrow=True. For UNLOAD results, that route fails.

2. Chunked UNLOAD reads with PolarsCursor reject pl.read_parquet()-only options

With the chunksize option and read options given to execute(), the result is scanned with pl.scan_parquet(location, **kwargs) (pyathena/polars/result_set.py:1041). The options are passed through unchanged. scan_parquet() takes no columns, use_pyarrow, pyarrow_options, or memory_map in Polars 1.39 or 2.0, so a whole read that works fails when chunksize is added:

cursor = connect(cursor_class=PolarsCursor).cursor(unload=True, chunksize=1000)
cursor.execute("SELECT ...", columns=["ts"])
# OperationalError from TypeError: scan_parquet() got an unexpected keyword argument 'columns'

Whole and chunked reads should behave the same under a user option (#1212, contract item 8). One way is to translate the options that scan_parquet() can express (columns as a select()). Options that it cannot express would need to be rejected, or the read would need to stay whole.

3. A failure while reading the data manifest body escapes OperationalError

AthenaResultSet._read_data_manifest() (pyathena/result_set.py:893) wraps get_object() in OperationalError, but it reads and decodes response["Body"] in the else: branch, outside the handler. If the body read fails after the response arrives, for example with a botocore ReadTimeoutError, the raw exception escapes from execute() of PandasCursor, ArrowCursor, and PolarsCursor with unload=True. Applications that catch PyAthena's DB API exceptions miss it. The get_object() call is retried, but the body read is not.

4. test_read_csv_with_pyarrow_matches_pandas fails at the minimum dependency versions

With pyarrow 22.0.0, Polars 1.39.0, and pandas 3.0.0 (the floors in pyproject.toml) on Python 3.13, 5 cases fail on master 9e2bf6e: types, dtype, parse_dates, dtype_of_date_column, and tab_separated_extra_fields, each with False. They fail on dtype attributes: object against StringDtype(na_value=nan) for 4 columns, and the tm and iv columns.

uv run --with pyarrow==22.0.0 --with polars==1.39.0 --with pandas==3.0.0 \
  pytest --noconftest tests/pyathena/pandas/test_result_set.py -k test_read_csv_with_pyarrow_matches_pandas

CI tests the locked versions only, so it does not run this combination. Either the CSV path or the test differs from pandas at pandas 3.0.0. Which one needs checking, and then either the code or the floor needs adjusting.

Decision (2026-10-18, revised)

An audit on Polars 1.39.0, 1.44.2, and 2.0.0 against live Athena found these additional cells:

  • With Polars 2.0, use_pyarrow=True on an UNLOAD result with timestamps fails with ValueError: schema cannot be used with use_pyarrow=True, because the result set adds schema= to read INT96 at microseconds.
  • Polars' pyarrow route cannot read a list of remote paths either (TypeError on a list of opened remote files), so passing the data manifest's paths alone does not fix item 1.
  • Chunked CSV reads have the same defect as item 2: scan_csv() takes no columns or use_pyarrow, and with Polars 1.x no n_threads, batch_size, or sample_size.
  • _read_s3_file_as_text() of the Spark cursors (get_std_out() / get_std_error()) has the defect of item 3, does not wrap a failed get_object() in OperationalError, and AioSparkCursor reads the body on the event loop.

Decided:

  1. Not fixed in PyAthena; reported upstream as read_parquet(use_pyarrow=True) cannot read a remote directory or a list of remote files pola-rs/polars#30038. The root cause is in Polars: prepare_file_arg(use_pyarrow=True) allows a local directory but opens a remote string path with fsspec.open(), which treats a prefix as one file, and _read_parquet_with_pyarrow() passes a list of opened remote files to pyarrow.parquet.read_table(), which rejects it. Reproduced without S3 on memory:// with Polars 1.44.2 and 2.0.0. PyAthena passes the option through and wraps the error in OperationalError, as before. The schema= conflict with Polars 2.0 becomes reachable only after the upstream fix.
  2. Not changed. The cursor fetches the query result; it does not emulate the semantics of Polars read options (columns, use_pyarrow, ...) for chunked scans. With the chunksize option, the options given to execute() go to scan_parquet() / scan_csv() unchanged, as the documentation says, and an option they do not take fails with OperationalError. The first PR's emulation and its fetch-row mapping were reverted.
  3. The body reads of the data manifest and the Spark output files are wrapped in OperationalError together with get_object(); AioSparkCursor reads the body in a thread. (Retrying the body read with get_object() was dropped: retry_api_call() retries only on the AWS error codes of RetryConfig, which a failed body read does not have.)
  4. Item 4 is fixed in the same PR: it is a pandas 3.0.0 bug (its PyArrow engine ignores future.infer_string=False; fixed in 3.0.1), so the test skips the pandas reference comparison there.

Closing PR: #1278 (replaces #1276, closed).

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