Read CSV results through fsspec for Polars 2.0

docs/polars.md

@@ -397,6 +397,7 @@ df = cursor.execute("SELECT * FROM many_rows").as_polars()
 They apply to the reads through PyAthena's filesystem: CSV results read without the chunksize option, which use `"bytes"` and 5 MiB blocks by default, and the Parquet files of UNLOAD results read with the chunksize option when pyarrow is installed and `execute()` is given no Polars read options.
 Of the latter, a file of up to 32 MiB is read at once, and a larger one is read with the given `cache_type`, or without it with the `"none"` cache, which requests only the ranges that pyarrow reads.
 With `storage_options` given to `execute()`, the given options replace PyAthena's S3 settings, and `block_size` and `cache_type` do not apply.
+For a CSV result read without the chunksize option, the options are passed to `fsspec.open()`, which uses PyAthena's S3FileSystem; for the other reads, they are passed to the Polars read function, which uses its native S3 access by default and fsspec with `use_pyarrow=True`.
 
 (async-polars-cursor)=
 

pyathena/polars/result_set.py

@@ -23,6 +23,8 @@
     cast,
 )
 
+import fsspec
+
 from pyathena import OperationalError
 from pyathena.arrow.reader import ParquetChunkReader
 from pyathena.converter import Converter
@@ -346,10 +348,10 @@ def _read_kwargs(
 
     @property
     def _csv_storage_options(self) -> dict[str, Any]:
-        """Get storage options for Polars CSV reading via fsspec.
+        """Get storage options for reading a CSV result file via fsspec.
 
-        Polars read_csv uses fsspec for cloud storage access, which works
-        with PyAthena's registered S3FileSystem.
+        ``_read_csv()`` opens the file with ``fsspec.open()``, which uses
+        PyAthena's registered S3FileSystem, and passes its contents to Polars.
 
         Returns:
             Dictionary with fsspec-compatible options for S3 access.
@@ -592,17 +594,22 @@ def _read_csv(self) -> pl.DataFrame:
         if "schema_overrides" not in self._kwargs:
             # 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.
+        storage_options = kwargs.pop("storage_options")
         source: str | bytes
         if self.output_location:
-            source = location = self.output_location
+            location = self.output_location
         else:
             source = self._fetch_all_rows_as_csv()
             if not source:
                 return pl.DataFrame()
             location = "the GetQueryResults rows"
-            del kwargs["storage_options"]
 
         try:
+            if self.output_location:
+                with fsspec.open(self.output_location, "rb", **(storage_options or {})) as file:
+                    source = file.read()
             try:
                 df = pl.read_csv(source, **kwargs)
             except pl.exceptions.ComputeError:

pyathena/polars/util.py

@@ -104,12 +104,18 @@ def get_athena_type(dtype: Any) -> tuple[str, int, int]:
         if base_dtype == polars_type or dtype == polars_type:
             return athena_info
 
-    # Handle parameterized types that didn't match above
-    dtype_str = str(dtype).lower()
-    if "list" in dtype_str:
+    # Nested types are matched by their outer type, not by the types they contain.
+    if base_dtype == pl.List:
         return ("array", 0, 0)
-    if "struct" in dtype_str:
+    if base_dtype == pl.Struct:
         return ("row", 0, 0)
+    # Polars 2.0 reads a Parquet map as Map; Polars 1.x reads it as a list of structs.
+    map_dtype = getattr(pl, "Map", None)
+    if map_dtype is not None and base_dtype == map_dtype:
+        return ("map", 0, 0)
+
+    # Handle parameterized types that didn't match above
+    dtype_str = str(dtype).lower()
     if "decimal" in dtype_str:
         # Extract precision and scale from Decimal type if available
         if hasattr(dtype, "precision") and hasattr(dtype, "scale"):

pyproject.toml

@@ -92,6 +92,10 @@ dev = [
     "pandas>=3.0.0",
     "pyarrow>=22.0.0",
     "polars>=1.39.0",
+    # Pull requests test the newest Python with Polars 2, and the older Pythons keep
+    # the newest Polars 1.x release.
+    "polars>=2.0.0; python_version >= '3.14'",
+    "polars<2.0.0; python_version < '3.14'",
     "Jinja2>=3.1.0",
     "mypy>=0.900",
     "pytest>=3.5",

tests/pyathena/arrow/test_cursor.py

@@ -626,6 +626,13 @@ def test_complex_unload_as_polars(self, arrow_cursor):
         ).as_polars()
         assert df.height == 1
         assert df.width == 16
+        # Polars 2.0 converts an Arrow map to Map; Polars 1.x to a list of key-value structs.
+        if hasattr(pl, "Map"):
+            map_dtype = pl.Map(pl.Int32, pl.Int32)
+            map_value = {1: 2, 3: 4}
+        else:
+            map_dtype = pl.List(pl.Struct([pl.Field("key", pl.Int32), pl.Field("value", pl.Int32)]))
+            map_value = [{"key": 1, "value": 2}, {"key": 3, "value": 4}]
         dtypes = tuple(df.dtypes)
         assert dtypes == (
             pl.Boolean,
@@ -641,7 +648,7 @@ def test_complex_unload_as_polars(self, arrow_cursor):
             pl.Date,
             pl.Binary,
             pl.List(pl.Int32),
-            pl.List(pl.Struct([pl.Field("key", pl.Int32), pl.Field("value", pl.Int32)])),
+            map_dtype,
             pl.Struct([pl.Field("a", pl.Int32), pl.Field("b", pl.Int32)]),
             pl.Decimal(precision=10, scale=1),
         )
@@ -661,7 +668,7 @@ def test_complex_unload_as_polars(self, arrow_cursor):
                 "col_date": datetime(2017, 1, 2).date(),
                 "col_binary": b"123",
                 "col_array": [1, 2],
-                "col_map": [{"key": 1, "value": 2}, {"key": 3, "value": 4}],
+                "col_map": map_value,
                 "col_struct": {"a": 1, "b": 2},
                 "col_decimal": Decimal("0.1"),
             }

tests/pyathena/polars/test_cursor.py

@@ -768,6 +768,29 @@ def test_description_with_chunksize_unload(self, polars_cursor, query, expected)
         assert description == ([("a", "integer"), ("b", "varchar")] if expected else [])
         assert polars_cursor.fetchall() == expected
 
+    @pytest.mark.parametrize(
+        "polars_cursor",
+        [
+            {"cursor_kwargs": {"unload": True}},
+            {"cursor_kwargs": {"unload": True, "chunksize": 5}},
+        ],
+        indirect=["polars_cursor"],
+    )
+    def test_nested_type_description_unload(self, polars_cursor):
+        """UNLOAD results describe nested Parquet columns by their outer type."""
+        polars_cursor.execute(
+            """
+            SELECT
+              ARRAY[1, 2] AS a
+              ,CAST(ROW(ARRAY[1]) AS ROW(x ARRAY(INTEGER))) AS r
+              ,MAP(ARRAY['k'], ARRAY[1]) AS m
+            """
+        )
+        # Polars 1.x reads a Parquet map as a list of key-value structs.
+        map_type = "map" if hasattr(pl, "Map") else "array"
+        description = [d[:2] for d in polars_cursor.description]
+        assert description == [("a", "array"), ("r", "row"), ("m", map_type)]
+
     @pytest.mark.parametrize(
         "polars_cursor",
         [
@@ -922,9 +945,12 @@ def test_new_columns_renaming_first_columns(self, polars_cursor, chunksize):
         assert polars_cursor.fetchall() == [("001", 2)]
 
     def test_new_columns_with_schema_overrides(self, polars_cursor):
-        """schema_overrides given with new_columns are keyed by the new names, as in Polars."""
+        """schema_overrides given with new_columns are keyed by the new names, as in Polars.
+
+        Polars 2.0 rejects new_columns that do not name every column of the file.
+        """
         polars_cursor.execute(
-            "SELECT '001' AS x, 2 AS y", new_columns=["z"], schema_overrides={"z": pl.String}
+            "SELECT '001' AS x, 2 AS y", new_columns=["z", "w"], schema_overrides={"z": pl.String}
         )
         assert polars_cursor.fetchall() == [("001", 2)]
 

tests/pyathena/polars/test_result_set.py

@@ -117,11 +117,13 @@ def test_csv_read_kwargs_replace_defaults(self, tmp_path, reader):
         """Read arguments given to execute() replace the ones the result set chooses."""
         path = tmp_path / "result.csv"
         path.write_text("1;x\n2;y\n")
+        # Polars names headerless columns from column_1 before 2.0 and from column_0 since.
+        first, second = pl.read_csv(b"1;x\n", separator=";", has_header=False).columns
         result_set = _chunked_result_set()
         result_set._kwargs = {
             "separator": ";",
             "has_header": False,
-            "schema_overrides": {"column_1": pl.Utf8},
+            "schema_overrides": {first: pl.Utf8},
         }
         with (
             patch.object(
@@ -152,15 +154,16 @@ def test_csv_read_kwargs_replace_defaults(self, tmp_path, reader):
         ):
             result = getattr(result_set, reader)()
             df = result if isinstance(result, pl.DataFrame) else pl.concat(list(result))
-        assert df.to_dict(as_series=False) == {"column_1": ["1", "2"], "column_2": ["x", "y"]}
+        assert df.to_dict(as_series=False) == {first: ["1", "2"], second: ["x", "y"]}
 
     @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"])
     def test_txt_new_columns_with_schema_overrides(self, tmp_path, reader):
         """schema_overrides given with new_columns reach Polars with them for a .txt file."""
         path = tmp_path / "result.txt"
         path.write_text("001\tx\n")
         result_set = _chunked_result_set()
-        result_set._kwargs = {"new_columns": ["z"], "schema_overrides": {"z": pl.Utf8}}
+        # Polars 2.0 rejects new_columns that do not name every column of the file.
+        result_set._kwargs = {"new_columns": ["z", "w"], "schema_overrides": {"z": pl.Utf8}}
         with (
             patch.object(
                 AthenaPolarsResultSet,
@@ -188,12 +191,11 @@ def test_txt_new_columns_with_schema_overrides(self, tmp_path, reader):
         ):
             result = getattr(result_set, reader)()
             df = result if isinstance(result, pl.DataFrame) else pl.concat(list(result))
-        assert df.to_dict(as_series=False) == {"z": ["001"], "b": ["x"]}
+        assert df.to_dict(as_series=False) == {"z": ["001"], "w": ["x"]}
 
     @pytest.mark.parametrize(
         ("reader", "function"),
         [
-            ("_read_csv", "read_csv"),
             ("_iter_csv_chunks", "scan_csv"),
             ("_read_parquet", "read_parquet"),
             ("_iter_parquet_chunks", "scan_parquet"),
@@ -237,6 +239,34 @@ def test_storage_options_replace_defaults(self, reader, function):
                 list(result)
         assert read.call_args.kwargs["storage_options"] == {"anon": True}
 
+    def test_csv_storage_options_replace_defaults(self):
+        """storage_options given to execute() open a CSV result through fsspec, not Polars."""
+        result_set = _chunked_result_set()
+        result_set._kwargs = {"storage_options": {"anon": True}}
+        with (
+            patch.object(
+                AthenaPolarsResultSet,
+                "output_location",
+                new_callable=PropertyMock,
+                return_value="s3://bucket/result.csv",
+            ),
+            patch.object(
+                AthenaPolarsResultSet, "_csv_dtypes", new_callable=PropertyMock, return_value={}
+            ),
+            patch.object(
+                AthenaPolarsResultSet,
+                "_csv_storage_options",
+                new_callable=PropertyMock,
+                side_effect=AssertionError("replaced storage options were computed"),
+            ),
+            patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True),
+            patch("pyathena.polars.result_set.fsspec.open") as open_,
+        ):
+            open_.return_value.__enter__.return_value.read.return_value = b"a\n1\n"
+            df = result_set._read_csv()
+        open_.assert_called_once_with("s3://bucket/result.csv", "rb", anon=True)
+        assert df.to_dict(as_series=False) == {"a": [1]}
+
     @pytest.mark.parametrize(
         ("kwargs", "expected"),
         [

tests/pyathena/polars/test_util.py

@@ -0,0 +1,31 @@
+# Copyright 2026 The PyAthena authors
+#
+# Licensed under the MIT License.
+# See LICENSE or https://opensource.org/licenses/MIT.
+#
+# SPDX-License-Identifier: MIT
+
+import polars as pl
+import pytest
+
+from pyathena.polars.util import get_athena_type
+
+
+@pytest.mark.parametrize(
+    ("dtype", "expected"),
+    [
+        (pl.List(pl.Int32), ("array", 0, 0)),
+        (pl.List(pl.Struct({"x": pl.Int32})), ("array", 0, 0)),
+        (pl.Struct({"x": pl.Int32}), ("row", 0, 0)),
+        (pl.Struct({"x": pl.List(pl.Int32)}), ("row", 0, 0)),
+    ],
+)
+def test_get_athena_type_nested(dtype, expected):
+    assert get_athena_type(dtype) == expected
+
+
+@pytest.mark.skipif(not hasattr(pl, "Map"), reason="Polars 1.x has no Map type")
+def test_get_athena_type_map():
+    assert get_athena_type(pl.Map(pl.String, pl.Int32)) == ("map", 0, 0)
+    assert get_athena_type(pl.Map(pl.String, pl.List(pl.Int32))) == ("map", 0, 0)
+    assert get_athena_type(pl.List(pl.Map(pl.String, pl.Int32))) == ("array", 0, 0)