Add AioCursorBase, AioS3FSCursor, and AioSparkCursor for Phase 3

Phase 3 of native asyncio cursor implementation:

Part 1 - Boilerplate deduplication:
- Extract shared properties, lifecycle methods, sync fetch, and async
  protocol into AioCursorBase (aio/base.py)
- Refactor AioCursor, AioPandasCursor, AioArrowCursor, AioPolarsCursor
  to extend AioCursorBase, reducing ~520 lines of duplicated code

Part 2 - AioS3FSCursor:
- Lightweight async CSV cursor using S3FileSystem
- Async fetch methods (via asyncio.to_thread) since S3FS uses lazy
  streaming from S3

Part 3 - AioSparkCursor:
- AioSparkBaseCursor overrides post-init I/O with async equivalents
  (poll, cancel, terminate_session, read_s3_file)
- AioSparkCursor for executing PySpark code asynchronously
- Session init stays sync (wrapped in asyncio.to_thread at creation)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: laughingman7743

pyathena/aio/arrow/cursor.py

@@ -3,9 +3,9 @@
 
 import asyncio
 import logging
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union, cast
 
-from pyathena.aio.common import AioBaseCursor
+from pyathena.aio.base import AioCursorBase
 from pyathena.arrow.converter import (
     DefaultArrowTypeConverter,
     DefaultArrowUnloadTypeConverter,
@@ -14,7 +14,6 @@
 from pyathena.common import CursorIterator
 from pyathena.error import OperationalError, ProgrammingError
 from pyathena.model import AthenaCompression, AthenaFileFormat, AthenaQueryExecution
-from pyathena.result_set import WithResultSet
 
 if TYPE_CHECKING:
     import polars as pl
@@ -23,7 +22,7 @@
 _logger = logging.getLogger(__name__)  # type: ignore
 
 
-class AioArrowCursor(AioBaseCursor, CursorIterator, WithResultSet):
+class AioArrowCursor(AioCursorBase):
     """Native asyncio cursor that returns results as Apache Arrow Tables.
 
     Uses ``asyncio.to_thread()`` to create the result set off the event loop.
@@ -66,13 +65,12 @@ def __init__(
             kill_on_interrupt=kill_on_interrupt,
             result_reuse_enable=result_reuse_enable,
             result_reuse_minutes=result_reuse_minutes,
+            on_start_query_execution=on_start_query_execution,
             **kwargs,
         )
         self._unload = unload
-        self._on_start_query_execution = on_start_query_execution
         self._connect_timeout = connect_timeout
         self._request_timeout = request_timeout
-        self._query_id: Optional[str] = None
         self._result_set: Optional[AthenaArrowResultSet] = None
 
     @staticmethod
@@ -83,45 +81,6 @@ def get_default_converter(
             return DefaultArrowUnloadTypeConverter()
         return DefaultArrowTypeConverter()
 
-    @property
-    def arraysize(self) -> int:
-        return self._arraysize
-
-    @arraysize.setter
-    def arraysize(self, value: int) -> None:
-        if value <= 0:
-            raise ProgrammingError("arraysize must be a positive integer value.")
-        self._arraysize = value
-
-    @property  # type: ignore
-    def result_set(self) -> Optional[AthenaArrowResultSet]:
-        return self._result_set
-
-    @result_set.setter
-    def result_set(self, val) -> None:
-        self._result_set = val
-
-    @property
-    def query_id(self) -> Optional[str]:
-        return self._query_id
-
-    @query_id.setter
-    def query_id(self, val) -> None:
-        self._query_id = val
-
-    @property
-    def rownumber(self) -> Optional[int]:
-        return self.result_set.rownumber if self.result_set else None
-
-    @property
-    def rowcount(self) -> int:
-        return self.result_set.rowcount if self.result_set else -1
-
-    def close(self) -> None:
-        """Close the cursor and release associated resources."""
-        if self.result_set and not self.result_set.is_closed:
-            self.result_set.close()
-
     async def execute(  # type: ignore[override]
         self,
         operation: str,
@@ -202,84 +161,6 @@ async def execute(  # type: ignore[override]
             raise OperationalError(query_execution.state_change_reason)
         return self
 
-    async def executemany(  # type: ignore[override]
-        self,
-        operation: str,
-        seq_of_parameters: List[Optional[Union[Dict[str, Any], List[str]]]],
-        **kwargs,
-    ) -> None:
-        """Execute a SQL query multiple times with different parameters.
-
-        Args:
-            operation: SQL query string to execute.
-            seq_of_parameters: Sequence of parameter sets, one per execution.
-            **kwargs: Additional keyword arguments passed to each ``execute()``.
-        """
-        for parameters in seq_of_parameters:
-            await self.execute(operation, parameters, **kwargs)
-        self._reset_state()
-
-    async def cancel(self) -> None:
-        """Cancel the currently executing query.
-
-        Raises:
-            ProgrammingError: If no query is currently executing.
-        """
-        if not self.query_id:
-            raise ProgrammingError("QueryExecutionId is none or empty.")
-        await self._cancel(self.query_id)
-
-    def fetchone(
-        self,
-    ) -> Optional[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
-        """Fetch the next row of the result set.
-
-        Returns:
-            A tuple representing the next row, or None if no more rows.
-
-        Raises:
-            ProgrammingError: If no result set is available.
-        """
-        if not self.has_result_set:
-            raise ProgrammingError("No result set.")
-        result_set = cast(AthenaArrowResultSet, self.result_set)
-        return result_set.fetchone()
-
-    def fetchmany(
-        self, size: Optional[int] = None
-    ) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
-        """Fetch multiple rows from the result set.
-
-        Args:
-            size: Maximum number of rows to fetch. Defaults to arraysize.
-
-        Returns:
-            List of tuples representing the fetched rows.
-
-        Raises:
-            ProgrammingError: If no result set is available.
-        """
-        if not self.has_result_set:
-            raise ProgrammingError("No result set.")
-        result_set = cast(AthenaArrowResultSet, self.result_set)
-        return result_set.fetchmany(size)
-
-    def fetchall(
-        self,
-    ) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
-        """Fetch all remaining rows from the result set.
-
-        Returns:
-            List of tuples representing all remaining rows.
-
-        Raises:
-            ProgrammingError: If no result set is available.
-        """
-        if not self.has_result_set:
-            raise ProgrammingError("No result set.")
-        result_set = cast(AthenaArrowResultSet, self.result_set)
-        return result_set.fetchall()
-
     def as_arrow(self) -> "Table":
         """Return query results as an Apache Arrow Table.
 
@@ -301,18 +182,3 @@ def as_polars(self) -> "pl.DataFrame":
             raise ProgrammingError("No result set.")
         result_set = cast(AthenaArrowResultSet, self.result_set)
         return result_set.as_polars()
-
-    def __aiter__(self):
-        return self
-
-    async def __anext__(self):
-        row = self.fetchone()
-        if row is None:
-            raise StopAsyncIteration
-        return row
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        self.close()

pyathena/aio/base.py

@@ -0,0 +1,158 @@
+# -*- coding: utf-8 -*-
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
+
+from pyathena.aio.common import AioBaseCursor
+from pyathena.common import CursorIterator
+from pyathena.error import ProgrammingError
+from pyathena.result_set import AthenaResultSet, WithResultSet
+
+
+class AioCursorBase(AioBaseCursor, CursorIterator, WithResultSet):
+    """Base class for native asyncio SQL cursors.
+
+    Provides shared properties, lifecycle methods, and default sync fetch
+    (for cursors whose result sets load all data eagerly in ``__init__``).
+    Subclasses override ``execute()`` and optionally ``__init__`` and
+    format-specific helpers.
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._query_id: Optional[str] = None
+        self._result_set: Optional[AthenaResultSet] = None
+        self._on_start_query_execution = kwargs.get("on_start_query_execution")
+
+    @property
+    def arraysize(self) -> int:
+        return self._arraysize
+
+    @arraysize.setter
+    def arraysize(self, value: int) -> None:
+        if value <= 0:
+            raise ProgrammingError("arraysize must be a positive integer value.")
+        self._arraysize = value
+
+    @property  # type: ignore
+    def result_set(self) -> Optional[AthenaResultSet]:
+        return self._result_set
+
+    @result_set.setter
+    def result_set(self, val) -> None:
+        self._result_set = val
+
+    @property
+    def query_id(self) -> Optional[str]:
+        return self._query_id
+
+    @query_id.setter
+    def query_id(self, val) -> None:
+        self._query_id = val
+
+    @property
+    def rownumber(self) -> Optional[int]:
+        return self.result_set.rownumber if self.result_set else None
+
+    @property
+    def rowcount(self) -> int:
+        return self.result_set.rowcount if self.result_set else -1
+
+    def close(self) -> None:
+        """Close the cursor and release associated resources."""
+        if self.result_set and not self.result_set.is_closed:
+            self.result_set.close()
+
+    async def executemany(  # type: ignore[override]
+        self,
+        operation: str,
+        seq_of_parameters: List[Optional[Union[Dict[str, Any], List[str]]]],
+        **kwargs,
+    ) -> None:
+        """Execute a SQL query multiple times with different parameters.
+
+        Args:
+            operation: SQL query string to execute.
+            seq_of_parameters: Sequence of parameter sets, one per execution.
+            **kwargs: Additional keyword arguments passed to each ``execute()``.
+        """
+        for parameters in seq_of_parameters:
+            await self.execute(operation, parameters, **kwargs)
+        # Operations that have result sets are not allowed with executemany.
+        self._reset_state()
+
+    async def cancel(self) -> None:
+        """Cancel the currently executing query.
+
+        Raises:
+            ProgrammingError: If no query is currently executing.
+        """
+        if not self.query_id:
+            raise ProgrammingError("QueryExecutionId is none or empty.")
+        await self._cancel(self.query_id)
+
+    def fetchone(
+        self,
+    ) -> Optional[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        """Fetch the next row of the result set.
+
+        Returns:
+            A tuple representing the next row, or None if no more rows.
+
+        Raises:
+            ProgrammingError: If no result set is available.
+        """
+        if not self.has_result_set:
+            raise ProgrammingError("No result set.")
+        result_set = cast(AthenaResultSet, self.result_set)
+        return result_set.fetchone()
+
+    def fetchmany(
+        self, size: Optional[int] = None
+    ) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        """Fetch multiple rows from the result set.
+
+        Args:
+            size: Maximum number of rows to fetch. Defaults to arraysize.
+
+        Returns:
+            List of tuples representing the fetched rows.
+
+        Raises:
+            ProgrammingError: If no result set is available.
+        """
+        if not self.has_result_set:
+            raise ProgrammingError("No result set.")
+        result_set = cast(AthenaResultSet, self.result_set)
+        return result_set.fetchmany(size)
+
+    def fetchall(
+        self,
+    ) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        """Fetch all remaining rows from the result set.
+
+        Returns:
+            List of tuples representing all remaining rows.
+
+        Raises:
+            ProgrammingError: If no result set is available.
+        """
+        if not self.has_result_set:
+            raise ProgrammingError("No result set.")
+        result_set = cast(AthenaResultSet, self.result_set)
+        return result_set.fetchall()
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        row = self.fetchone()
+        if row is None:
+            raise StopAsyncIteration
+        return row
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.close()