diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 957f847b3..fdd672355 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -16,6 +16,8 @@ import threading from collections import abc from collections.abc import Callable, Iterable, Iterator +from contextlib import AbstractContextManager, nullcontext +from functools import partial from io import BytesIO, StringIO from itertools import islice, repeat from multiprocessing import cpu_count @@ -607,7 +609,7 @@ def _csv_storage_options(self) -> dict[str, Any]: """Get storage options for reading a CSV result file via fsspec. ``_read_csv()`` opens the file with ``fsspec.open()``, which uses - PyAthena's registered S3FileSystem, and passes its contents to Polars. + PyAthena's registered S3FileSystem, and passes the open file to Polars. Returns: Dictionary with fsspec-compatible options for S3 access. @@ -852,9 +854,9 @@ def _read_csv(self) -> pl.DataFrame: # Renamed after reading, so that Polars matches the types to the header. kwargs.pop("new_columns", None) # Polars 2.0 reads a path with its native object store, which does not take - # fsspec options, so the result file is read through fsspec here. + # fsspec options, so the result file is opened through fsspec here and given + # to Polars as a file, which Polars 1.x does not copy as it does bytes. storage_options = kwargs.pop("storage_options") - source: str | bytes if self.output_location: location = self.output_location else: @@ -864,15 +866,22 @@ def _read_csv(self) -> pl.DataFrame: location = "the GetQueryResults rows" try: - if self.output_location: - with fsspec.open(self.output_location, "rb", **(storage_options or {})) as file: - source = file.read() + # Each read calls open_csv() to open the file again and closes it after + # reading; reading the same file again after seek(0) makes fsspec's cache + # hold a second copy. It is created and called inside the try, so that + # errors such as invalid storage options raise OperationalError. + open_csv: Callable[[], AbstractContextManager[Any]] = ( + partial(fsspec.open, self.output_location, "rb", **(storage_options or {})) + if self.output_location + else partial(nullcontext, source) + ) schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) strict = not kwargs.get("ignore_errors", False) df: pl.DataFrame | None = None if strict or not timestamp_dtypes: try: - df = pl.read_csv(source, **kwargs) + with open_csv() as data: + df = pl.read_csv(data, **kwargs) except pl.exceptions.ComputeError: if not timestamp_dtypes: raise @@ -881,7 +890,9 @@ def _read_csv(self) -> pl.DataFrame: # failed, or at once with ignore_errors, with which Polars reads the # timestamps that it does not parse as null instead of failing. kwargs["schema_overrides"] = schema_overrides - df = _to_datetimes(pl.read_csv(source, **kwargs), timestamp_dtypes, strict=strict) + with open_csv() as data: + texts = pl.read_csv(data, **kwargs) + df = _to_datetimes(texts, timestamp_dtypes, strict=strict) if new_columns: df.columns = [*new_columns, *df.columns[len(new_columns) :]] return df diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index a5df9cf78..e744d8967 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -14,6 +14,7 @@ import polars as pl import pytest +from botocore.exceptions import ProfileNotFound from fsspec.implementations.local import LocalFileSystem from pyathena.converter import _to_date, _to_default @@ -276,6 +277,17 @@ def test_csv_chunks_raise_when_read_fails_partway(self, tmp_path): with pytest.raises(OperationalError): list(result_set.iter_chunks()) + @pytest.mark.parametrize( + ("storage_options", "cause"), + [({"profile": "pyathena-no-such-profile"}, ProfileNotFound), ("invalid", TypeError)], + ids=["file-system", "not-a-mapping"], + ) + def test_csv_read_raises_with_invalid_storage_options(self, storage_options, cause): + """A whole CSV read raises OperationalError when its storage options are invalid.""" + with pytest.raises(OperationalError) as e: + _local_result_set([("a", "integer")], storage_options=storage_options) + assert isinstance(e.value.__cause__, cause) + @pytest.mark.parametrize( "read_options", [{}, {"low_memory": True}],