Backport #1185 and #1169 to 3.x

pyathena/aio/sqlalchemy/base.py

@@ -10,6 +10,8 @@
 
 import pyathena
 from pyathena.aio.connection import AioConnection
+from pyathena.aio.cursor import AioCursor
+from pyathena.cursor import Cursor
 from pyathena.error import (
     DatabaseError,
     DataError,
@@ -29,6 +31,9 @@
     from sqlalchemy import URL
 
 
+_ASYNC_CURSOR_CLASSES: dict[Any, Any] = {Cursor: AioCursor}
+
+
 class AsyncAdapt_pyathena_cursor:
     """Wraps any async PyAthena cursor with a sync DBAPI interface.
 
@@ -139,8 +144,10 @@ def catalog_name(self) -> str | None:
     def schema_name(self) -> str | None:
         return self._connection.schema_name  # type: ignore[no-any-return]
 
-    def cursor(self) -> AsyncAdapt_pyathena_cursor:
-        raw_cursor = self._connection.cursor()
+    def cursor(self, cursor: Any = None) -> AsyncAdapt_pyathena_cursor:
+        # The shared dialect names a cursor class in its synchronous form; this
+        # connection can only drive the async counterpart.
+        raw_cursor = self._connection.cursor(_ASYNC_CURSOR_CLASSES.get(cursor, cursor))
         return AsyncAdapt_pyathena_cursor(raw_cursor)
 
     def close(self) -> None:

pyathena/polars/result_set.py

@@ -10,6 +10,8 @@
     cast,
 )
 
+import fsspec
+
 from pyathena import OperationalError
 from pyathena.converter import Converter
 from pyathena.error import ProgrammingError
@@ -268,10 +270,10 @@ def __init__(
 
     @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.
@@ -463,12 +465,15 @@ def _read_csv(self) -> pl.DataFrame:
         separator, has_header, new_columns = self._get_csv_params()
 
         try:
+            # 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.
+            with fsspec.open(self.output_location, "rb", **self._csv_storage_options) as file:
+                source = file.read()
             df = pl.read_csv(
-                self.output_location,
+                source,
                 separator=separator,
                 has_header=has_header,
                 schema_overrides=self.dtypes,
-                storage_options=self._csv_storage_options,
                 **self._kwargs,
             )
             if new_columns:

pyathena/polars/util.py

@@ -97,12 +97,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"):

pyathena/sqlalchemy/base.py

@@ -23,6 +23,7 @@
 )
 
 import pyathena
+from pyathena.cursor import Cursor
 from pyathena.sqlalchemy.compiler import (
     AthenaDDLCompiler,
     AthenaStatementCompiler,
@@ -52,6 +53,7 @@
         PoolProxiedConnection,
     )
     from sqlalchemy.engine.interfaces import (
+        DBAPIConnection,
         ReflectedForeignKeyConstraint,
         ReflectedIndex,
         ReflectedPrimaryKeyConstraint,
@@ -85,6 +87,35 @@
 }
 
 
+class _PingConnection:
+    """A DBAPI connection whose ``cursor()`` opens an API cursor.
+
+    ``AthenaDialect.do_ping()`` passes it to ``DefaultDialect.do_ping()``, so the ping
+    does not run through the connection's default cursor class. Other attributes are
+    those of the pinged connection, which ``is_disconnect()`` receives before
+    SQLAlchemy 2.0.5.
+    """
+
+    def __init__(self, connection: Any) -> None:
+        """Wrap the pinged connection.
+
+        Args:
+            connection: The pinged DBAPI connection, or the async adapter around one.
+        """
+        self._connection = connection
+
+    def cursor(self) -> Any:
+        """Open an API cursor on the pinged connection.
+
+        Returns:
+            The cursor, wrapped by the async adapter for an async connection.
+        """
+        return self._connection.cursor(Cursor)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._connection, name)
+
+
 class AthenaDialect(DefaultDialect):
     """SQLAlchemy dialect for Amazon Athena.
 
@@ -284,6 +315,14 @@ def _get_table(self, connection, table_name: str, schema: str | None = None, **k
                     raise exc.NoSuchTableError(table_name) from e
                 raise
 
+    def do_ping(self, dbapi_connection: DBAPIConnection) -> bool:
+        # When unload is enabled in the default cursor, SELECT 1 is wrapped in UNLOAD
+        # which Athena rejects with MISSING_COLUMN_NAME because the column is unnamed.
+        # Using the plain cursor here performs the SELECT without the UNLOAD.
+        # The ping still runs in DefaultDialect.do_ping(), which before SQLAlchemy 2.0.5
+        # is also where a failed ping is checked for a disconnect.
+        return super().do_ping(cast("DBAPIConnection", _PingConnection(dbapi_connection)))
+
     @reflection.cache
     def _get_tables(self, connection, schema: str | None = None, **kw):
         raw_connection = self._raw_connection(connection)

pyproject.toml

@@ -75,6 +75,10 @@ dev = [
     "pyarrow>=10.0.0; python_version<'3.14'",
     "pyarrow>=22.0.0; python_version>='3.14'",
     "polars>=1.0.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/aio/sqlalchemy/test_base.py

@@ -7,6 +7,17 @@
 
 
 class TestAsyncSQLAlchemyAthena:
+    @pytest.mark.parametrize(
+        "async_engine",
+        [{"driver": "aiopandas", "unload": True}],
+        indirect=True,
+    )
+    async def test_ping_with_unload(self, async_engine):
+        _, conn = async_engine
+        assert await conn.run_sync(
+            lambda sync_conn: sync_conn.dialect.do_ping(sync_conn.connection.dbapi_connection)
+        )
+
     @pytest.mark.parametrize(
         "async_engine",
         [

tests/pyathena/arrow/test_cursor.py

@@ -609,6 +609,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,
@@ -624,7 +631,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),
         )
@@ -644,7 +651,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

@@ -635,6 +635,28 @@ def test_iterator_with_chunksize_unload(self, polars_cursor):
         rows = list(polars_cursor)
         assert len(rows) == 15
 
+    @pytest.mark.parametrize(
+        "polars_cursor",
+        [
+            {"cursor_kwargs": {"unload": True}},
+        ],
+        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",
         [

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)