Merge pull request #828 from pyathena-dev/fix/820-821-polars-floor

Require polars>=1.39.0 for chunked PolarsCursor reads

Author: laughingman7743

README.md

@@ -48,7 +48,7 @@ Extra packages:
 | AioSQLAlchemy | `pip install PyAthena[AioSQLAlchemy]`   | >=2.0.0  |
 | Pandas        | `pip install PyAthena[Pandas]`          | >=1.3.0  |
 | Arrow         | `pip install PyAthena[Arrow]`           | >=10.0.0 |
-| Polars        | `pip install PyAthena[Polars]`          | >=1.0.0  |
+| Polars        | `pip install PyAthena[Polars]`          | >=1.39.0 |
 
 ## Usage
 

benchmarks/README.md

@@ -227,8 +227,8 @@ API row conversion and native DataFrame/Table access appear as separate cases.
 Sampled RSS can miss short peaks, and the constructor suite's RSS includes its subsequent validation read.
 If the operating system denies access to per-thread CPU times, `thread_cpu_available` is false; thread counts and RSS are still recorded.
 Each trial process receives its own `POLARS_TEMP_DIR`, which the parent samples with RSS and removes after the trial.
-Polars lazy CSV scans of S3 objects download the whole object into a file cache in this directory before yielding batches; in the recorded runs, those files remained after the process exited.
-When the temporary directory is a tmpfs, as `/tmp` is on Amazon Linux 2023, this storage uses memory that RSS does not include.
+Polars before 1.39.0 downloaded the whole S3 object of a lazy CSV scan into a file cache in this directory before yielding batches and left the file after the process exited; PyAthena requires Polars 1.39.0 or later, whose scans do not write that cache.
+When the temporary directory is a tmpfs, as `/tmp` is on Amazon Linux 2023, files written there use memory that RSS does not include.
 
 | Capability | Treatment |
 | --- | --- |

benchmarks/pyathena_bench/runner.py

@@ -175,9 +175,10 @@ def directory_bytes(path: Path) -> int:
 def start_with_temp_dir(process: Any, temp_dir: Path) -> None:
     """Start a spawned trial process with its own Polars temporary directory.
 
-    Polars downloads cloud objects for lazy scans into a file cache under
-    ``POLARS_TEMP_DIR`` and keeps the files after the process exits. A
-    per-trial directory lets the parent measure and remove that storage.
+    Polars before 1.39.0 downloads cloud objects for lazy CSV scans into a
+    file cache under ``POLARS_TEMP_DIR`` and keeps the files after the process
+    exits. A per-trial directory lets the parent measure and remove any
+    storage Polars writes there.
 
     Args:
         process: Unstarted multiprocessing process; it inherits the environment.

docs/introduction.md

@@ -35,7 +35,7 @@ Extra packages:
 | AioSQLAlchemy | `pip install PyAthena[AioSQLAlchemy]`   | >=2.0.0  |
 | Pandas        | `pip install PyAthena[Pandas]`          | >=1.3.0  |
 | Arrow         | `pip install PyAthena[Arrow]`           | >=10.0.0 |
-| Polars        | `pip install PyAthena[Polars]`          | >=1.0.0  |
+| Polars        | `pip install PyAthena[Polars]`          | >=1.39.0 |
 
 (features)=
 

docs/polars.md

@@ -11,8 +11,9 @@ This cursor downloads the CSV file after executing the query and loads it into a
 Performance is better than fetching data with Cursor.
 
 PolarsCursor uses [Polars](https://pola.rs/) native reading capabilities (`pl.read_csv`, `pl.read_parquet`) and
-does not require PyArrow as a dependency. PyAthena's own S3FileSystem (fsspec compatible)
-is used for S3 access, so s3fs is also not required.
+does not require PyArrow as a dependency. CSV results read without the chunksize option use
+PyAthena's own S3FileSystem (fsspec compatible), and other reads use Polars' native S3 access,
+so s3fs is also not required.
 
 You can use the PolarsCursor by specifying the `cursor_class`
 with the connect method or connection object.
@@ -310,8 +311,11 @@ for chunk in cursor.iter_chunks():
     print(f"Processed chunk with {chunk.height} rows")
 ```
 
-This method uses Polars' `scan_csv()` and `scan_parquet()` with `collect_batches()`
-for efficient lazy evaluation, minimizing memory usage when processing large datasets.
+This method uses Polars' `scan_csv()` and `scan_parquet()` with `collect_batches()`,
+which read the result from S3 without writing it to local storage.
+Memory usage depends on the size of each chunk and on how far Polars reads ahead, which `chunksize` does not limit.
+For CSV results, Polars limits the read-ahead by the number of CPU cores, so memory usage does not grow with the result size.
+If reading a chunk fails, iteration raises `OperationalError` instead of ending early.
 
 The chunked iteration also works with the unload option:
 

pyproject.toml

@@ -65,7 +65,7 @@ arrow = [
     "pyarrow>=22.0.0; python_version>='3.14'",
 ]
 polars = [
-    "polars>=1.0.0",
+    "polars>=1.39.0",
 ]
 
 [dependency-groups]
@@ -77,7 +77,7 @@ dev = [
     "numpy>=2.3.0; python_version>='3.14'",
     "pyarrow>=10.0.0; python_version<'3.14'",
     "pyarrow>=22.0.0; python_version>='3.14'",
-    "polars>=1.0.0",
+    "polars>=1.39.0",
     "Jinja2>=3.1.0",
     "mypy>=0.900",
     "pytest>=3.5",

tests/pyathena/polars/test_result_set.py

@@ -0,0 +1,83 @@
+# Copyright 2026 The PyAthena authors
+#
+# Licensed under the MIT License.
+# See LICENSE or https://opensource.org/licenses/MIT.
+#
+# SPDX-License-Identifier: MIT
+
+from unittest.mock import PropertyMock, patch
+
+import polars as pl
+import pytest
+
+from pyathena.error import OperationalError
+from pyathena.polars.result_set import AthenaPolarsResultSet
+
+_ROWS_BEFORE_FAILURE = 300_000
+
+
+def _chunked_result_set() -> AthenaPolarsResultSet:
+    """Create a chunked result set without calling Athena.
+
+    Returns:
+        An AthenaPolarsResultSet with only the attributes used by the chunk readers.
+    """
+    result_set = AthenaPolarsResultSet.__new__(AthenaPolarsResultSet)  # bypass __init__
+    result_set._chunksize = 10_000
+    result_set._kwargs = {}
+    return result_set
+
+
+class TestAthenaPolarsResultSet:
+    def test_iter_csv_chunks_raises_when_read_fails_partway(self, tmp_path):
+        """A CSV read that fails partway through the data raises instead of ending early."""
+        path = tmp_path / "result.csv"
+        path.write_text(
+            "a\n" + "".join(f"{i}\n" for i in range(_ROWS_BEFORE_FAILURE)) + "not-a-number\n"
+        )
+        result_set = _chunked_result_set()
+        with (
+            patch.object(
+                AthenaPolarsResultSet,
+                "output_location",
+                new_callable=PropertyMock,
+                return_value=str(path),
+            ),
+            patch.object(
+                AthenaPolarsResultSet,
+                "dtypes",
+                new_callable=PropertyMock,
+                return_value={"a": pl.Int64},
+            ),
+            patch.object(
+                AthenaPolarsResultSet,
+                "_parquet_storage_options",
+                new_callable=PropertyMock,
+                return_value={},
+            ),
+            patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True),
+            pytest.raises(OperationalError, match="not-a-number"),
+        ):
+            list(result_set._iter_csv_chunks())
+
+    def test_iter_parquet_chunks_raises_when_read_fails_partway(self, tmp_path):
+        """A Parquet read that fails partway through the data raises instead of ending early."""
+        pl.DataFrame({"a": range(_ROWS_BEFORE_FAILURE)}).write_parquet(
+            tmp_path / "0.parquet", row_group_size=10_000
+        )
+        valid = (tmp_path / "0.parquet").read_bytes()
+        # Keep the footer magic but corrupt the metadata so the second file fails to read.
+        (tmp_path / "1.parquet").write_bytes(valid[: len(valid) // 2] + b"\x00" * 64 + valid[-8:])
+        result_set = _chunked_result_set()
+        result_set._unload_location = f"{tmp_path}/"
+        with (
+            patch.object(
+                AthenaPolarsResultSet,
+                "_parquet_storage_options",
+                new_callable=PropertyMock,
+                return_value={},
+            ),
+            patch.object(AthenaPolarsResultSet, "_prepare_parquet_location", return_value=True),
+            pytest.raises(OperationalError),
+        ):
+            list(result_set._iter_parquet_chunks())