From 87313988cbe87cf331540801be727508d9765114 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 20 Oct 2026 00:29:01 +0900 Subject: [PATCH 1/6] Pass the open CSV result file to Polars instead of its bytes (GH-1307) Since GH-1169, an unchunked PolarsCursor CSV read passed the whole result file to pl.read_csv() as bytes. Polars 1.x wraps bytes in a BytesIO and calls getbuffer() on it in _check_empty(), which makes CPython copy the shared buffer, so the read held a second copy of the file. Polars reads an open file object with one read() call and wraps the returned bytes without copying, in 1.x and 2.0, which is also what Polars 1.x did with the path and fsspec storage options before GH-1169. The result set now passes the file opened with fsspec.open() and seeks back to its start before the second read that parses timestamps as text. Rows fetched with GetQueryResults are still passed as bytes. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 50 +++++++++++++++++++---------------- 1 file changed, 27 insertions(+), 23 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 957f847b3..aa42f1bbd 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -16,6 +16,7 @@ import threading from collections import abc from collections.abc import Callable, Iterable, Iterator +from contextlib import AbstractContextManager, nullcontext from io import BytesIO, StringIO from itertools import islice, repeat from multiprocessing import cpu_count @@ -607,7 +608,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,36 +853,39 @@ 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. Polars + # reads the open file into one buffer, while Polars 1.x copies bytes given to it. storage_options = kwargs.pop("storage_options") - source: str | bytes + source: AbstractContextManager[Any] if self.output_location: location = self.output_location + source = fsspec.open(self.output_location, "rb", **(storage_options or {})) else: - source = self._fetch_all_rows_as_csv() - if not source: + rows = self._fetch_all_rows_as_csv() + if not rows: return pl.DataFrame() + source = nullcontext(rows) location = "the GetQueryResults rows" try: - if self.output_location: - with fsspec.open(self.output_location, "rb", **(storage_options or {})) as file: - source = file.read() - 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) - except pl.exceptions.ComputeError: - if not timestamp_dtypes: - raise - if df is None: - # The timestamp columns are read as text, again after parsing them - # 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 source as data: + 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(data, **kwargs) + except pl.exceptions.ComputeError: + if not timestamp_dtypes: + raise + if df is None: + # The timestamp columns are read as text, again after parsing them + # 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 + if not isinstance(data, bytes): + data.seek(0) + df = _to_datetimes(pl.read_csv(data, **kwargs), timestamp_dtypes, strict=strict) if new_columns: df.columns = [*new_columns, *df.columns[len(new_columns) :]] return df From 818f0b856e9d6afa45b625032d14ea60ac8e7291 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 20 Oct 2026 00:45:16 +0900 Subject: [PATCH 2/6] Open the CSV result file inside the read's error handling (GH-1307) fsspec.open() creates the filesystem when it is called, so an error in creating it, such as a profile in the storage options that does not exist, escaped as the botocore error instead of OperationalError. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 11 ++++++----- tests/pyathena/polars/test_result_set.py | 7 +++++++ 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index aa42f1bbd..0d50355e6 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -16,7 +16,7 @@ import threading from collections import abc from collections.abc import Callable, Iterable, Iterator -from contextlib import AbstractContextManager, nullcontext +from contextlib import nullcontext from io import BytesIO, StringIO from itertools import islice, repeat from multiprocessing import cpu_count @@ -856,19 +856,20 @@ def _read_csv(self) -> pl.DataFrame: # fsspec options, so the result file is opened through fsspec here. Polars # reads the open file into one buffer, while Polars 1.x copies bytes given to it. storage_options = kwargs.pop("storage_options") - source: AbstractContextManager[Any] if self.output_location: location = self.output_location - source = fsspec.open(self.output_location, "rb", **(storage_options or {})) else: rows = self._fetch_all_rows_as_csv() if not rows: return pl.DataFrame() - source = nullcontext(rows) location = "the GetQueryResults rows" try: - with source as data: + with ( + fsspec.open(self.output_location, "rb", **(storage_options or {})) + if self.output_location + else nullcontext(rows) + ) as data: schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) strict = not kwargs.get("ignore_errors", False) df: pl.DataFrame | None = None diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index a5df9cf78..50a7a7515 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -276,6 +276,13 @@ def test_csv_chunks_raise_when_read_fails_partway(self, tmp_path): with pytest.raises(OperationalError): list(result_set.iter_chunks()) + def test_csv_read_raises_when_file_system_fails(self): + """A whole CSV read raises OperationalError when the filesystem cannot be created.""" + with pytest.raises(OperationalError): + _local_result_set( + [("a", "integer")], storage_options={"profile": "pyathena-no-such-profile"} + ) + @pytest.mark.parametrize( "read_options", [{}, {"low_memory": True}], From 94531dfb921f858ef99c9110c9d45cd433c24423 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 20 Oct 2026 00:52:56 +0900 Subject: [PATCH 3/6] Close the CSV result file before converting timestamps (GH-1307) With a cache type such as "all" or "readahead", the open file can keep the bytes it read until it is closed, so the file is now open only for the reads. The second read seeks to the start when a file was opened, and the filesystem test checks that the error comes from the profile. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 12 +++++++----- tests/pyathena/polars/test_result_set.py | 4 +++- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 0d50355e6..a48ba50c9 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -865,14 +865,14 @@ def _read_csv(self) -> pl.DataFrame: location = "the GetQueryResults rows" try: + schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) + strict = not kwargs.get("ignore_errors", False) + df: pl.DataFrame | None = None with ( fsspec.open(self.output_location, "rb", **(storage_options or {})) if self.output_location else nullcontext(rows) ) as data: - 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(data, **kwargs) @@ -884,9 +884,11 @@ 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 - if not isinstance(data, bytes): + if self.output_location: data.seek(0) - df = _to_datetimes(pl.read_csv(data, **kwargs), timestamp_dtypes, strict=strict) + texts = pl.read_csv(data, **kwargs) + if df is None: + 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 50a7a7515..8272330e6 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 @@ -278,10 +279,11 @@ def test_csv_chunks_raise_when_read_fails_partway(self, tmp_path): def test_csv_read_raises_when_file_system_fails(self): """A whole CSV read raises OperationalError when the filesystem cannot be created.""" - with pytest.raises(OperationalError): + with pytest.raises(OperationalError) as e: _local_result_set( [("a", "integer")], storage_options={"profile": "pyathena-no-such-profile"} ) + assert isinstance(e.value.__cause__, ProfileNotFound) @pytest.mark.parametrize( "read_options", From 0f76f074a4fd805552f7e8813f4f769704e8d1f7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 20 Oct 2026 01:00:12 +0900 Subject: [PATCH 4/6] Open the CSV result file again for the timestamp read (GH-1307) After seek(0), the second read of the same file missed fsspec's BytesCache, which then joined the fetched range onto its tail and held two copies of the file. Each read now opens the file, so the second one starts with an empty cache, and the file is closed before the timestamps are converted without seeking or a separate open block. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 48 ++++++++++++++++++----------------- 1 file changed, 25 insertions(+), 23 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index a48ba50c9..42ba75dde 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -16,7 +16,8 @@ import threading from collections import abc from collections.abc import Callable, Iterable, Iterator -from contextlib import nullcontext +from contextlib import AbstractContextManager, nullcontext +from functools import partial from io import BytesIO, StringIO from itertools import islice, repeat from multiprocessing import cpu_count @@ -853,41 +854,42 @@ 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 opened through fsspec here. Polars - # reads the open file into one buffer, while Polars 1.x copies bytes given to it. + # 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. Each + # read opens the file again, so that a second read starts without the cache + # of the first. storage_options = kwargs.pop("storage_options") + open_source: Callable[[], AbstractContextManager[Any]] if self.output_location: location = self.output_location + open_source = partial( + fsspec.open, self.output_location, "rb", **(storage_options or {}) + ) else: - rows = self._fetch_all_rows_as_csv() - if not rows: + source = self._fetch_all_rows_as_csv() + if not source: return pl.DataFrame() + open_source = partial(nullcontext, source) location = "the GetQueryResults rows" try: schema_overrides, timestamp_dtypes = self._get_csv_schema_overrides(has_header) strict = not kwargs.get("ignore_errors", False) df: pl.DataFrame | None = None - with ( - fsspec.open(self.output_location, "rb", **(storage_options or {})) - if self.output_location - else nullcontext(rows) - ) as data: - if strict or not timestamp_dtypes: - try: + if strict or not timestamp_dtypes: + try: + with open_source() as data: df = pl.read_csv(data, **kwargs) - except pl.exceptions.ComputeError: - if not timestamp_dtypes: - raise - if df is None: - # The timestamp columns are read as text, again after parsing them - # 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 - if self.output_location: - data.seek(0) - texts = pl.read_csv(data, **kwargs) + except pl.exceptions.ComputeError: + if not timestamp_dtypes: + raise if df is None: + # The timestamp columns are read as text, again after parsing them + # 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 + with open_source() 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) :]] From ba06c2eab2285fc865617087f49bda7f73248f69 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 21 Oct 2026 00:40:16 +0900 Subject: [PATCH 5/6] Create the CSV result file opener inside the read's error handling (GH-1307) Unpacking storage_options that are not a mapping raised TypeError outside the try, so it escaped without OperationalError. The opener is now created inside it, and the test covers both storage options that fail to create the filesystem and ones that are not a mapping. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 10 +++++----- tests/pyathena/polars/test_result_set.py | 15 +++++++++------ 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 42ba75dde..cb5eb7763 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -859,20 +859,20 @@ def _read_csv(self) -> pl.DataFrame: # read opens the file again, so that a second read starts without the cache # of the first. storage_options = kwargs.pop("storage_options") - open_source: Callable[[], AbstractContextManager[Any]] if self.output_location: location = self.output_location - open_source = partial( - fsspec.open, self.output_location, "rb", **(storage_options or {}) - ) else: source = self._fetch_all_rows_as_csv() if not source: return pl.DataFrame() - open_source = partial(nullcontext, source) location = "the GetQueryResults rows" try: + open_source: 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 diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 8272330e6..e744d8967 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -277,13 +277,16 @@ def test_csv_chunks_raise_when_read_fails_partway(self, tmp_path): with pytest.raises(OperationalError): list(result_set.iter_chunks()) - def test_csv_read_raises_when_file_system_fails(self): - """A whole CSV read raises OperationalError when the filesystem cannot be created.""" + @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={"profile": "pyathena-no-such-profile"} - ) - assert isinstance(e.value.__cause__, ProfileNotFound) + _local_result_set([("a", "integer")], storage_options=storage_options) + assert isinstance(e.value.__cause__, cause) @pytest.mark.parametrize( "read_options", From 7bf315c1b7e811fa84459c08502ebe04dceffe85 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 21 Oct 2026 01:15:01 +0900 Subject: [PATCH 6/6] Name the CSV opener open_csv and explain why it is a factory (GH-1307) The local factory was named open_source, which also reads as "open source". It is now open_csv, the name the Arrow result set uses for the same job, and the comment above it says why each read opens the file again and why it is created inside the try. Co-Authored-By: Claude Opus 5.5 --- pyathena/polars/result_set.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index cb5eb7763..fdd672355 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -855,9 +855,7 @@ def _read_csv(self) -> pl.DataFrame: 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 opened through fsspec here and given - # to Polars as a file, which Polars 1.x does not copy as it does bytes. Each - # read opens the file again, so that a second read starts without the cache - # of the first. + # to Polars as a file, which Polars 1.x does not copy as it does bytes. storage_options = kwargs.pop("storage_options") if self.output_location: location = self.output_location @@ -868,7 +866,11 @@ def _read_csv(self) -> pl.DataFrame: location = "the GetQueryResults rows" try: - open_source: Callable[[], AbstractContextManager[Any]] = ( + # 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) @@ -878,7 +880,7 @@ def _read_csv(self) -> pl.DataFrame: df: pl.DataFrame | None = None if strict or not timestamp_dtypes: try: - with open_source() as data: + with open_csv() as data: df = pl.read_csv(data, **kwargs) except pl.exceptions.ComputeError: if not timestamp_dtypes: @@ -888,7 +890,7 @@ 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 - with open_source() as data: + with open_csv() as data: texts = pl.read_csv(data, **kwargs) df = _to_datetimes(texts, timestamp_dtypes, strict=strict) if new_columns: