Pass the open CSV result file to Polars instead of its bytes (GH-1307)

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

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}],