From dda6aaed9e0d7edfa7346d2a2b4cc50a0d79289e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 6 Oct 2026 01:21:15 +0900 Subject: [PATCH 1/6] Separate cursor query and S3 read worker settings --- docs/aio.md | 2 + docs/arrow.md | 2 + docs/cursor.md | 45 +++++++ docs/pandas.md | 2 + docs/polars.md | 2 + pyathena/aio/arrow/cursor.py | 23 +++- pyathena/aio/pandas/cursor.py | 14 +- pyathena/aio/polars/cursor.py | 14 +- pyathena/arrow/async_cursor.py | 28 ++-- pyathena/arrow/cursor.py | 26 ++-- pyathena/arrow/result_set.py | 96 +++++++++---- pyathena/async_cursor.py | 2 +- pyathena/filesystem/s3.py | 11 +- pyathena/pandas/async_cursor.py | 12 +- pyathena/pandas/cursor.py | 14 +- pyathena/polars/async_cursor.py | 13 +- pyathena/polars/cursor.py | 14 +- pyathena/util.py | 24 ++++ tests/pyathena/aio/arrow/test_cursor.py | 31 ++++- tests/pyathena/aio/pandas/test_cursor.py | 14 +- tests/pyathena/aio/polars/test_cursor.py | 14 +- tests/pyathena/arrow/test_async_cursor.py | 35 ++++- tests/pyathena/arrow/test_cursor.py | 31 ++++- tests/pyathena/arrow/test_result_set.py | 149 ++++++++++++++++++++- tests/pyathena/filesystem/test_s3.py | 14 ++ tests/pyathena/pandas/test_async_cursor.py | 22 ++- tests/pyathena/pandas/test_cursor.py | 14 +- tests/pyathena/polars/test_async_cursor.py | 17 ++- tests/pyathena/polars/test_cursor.py | 14 +- tests/pyathena/test_util.py | 82 ++++++++++++ 30 files changed, 667 insertions(+), 114 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 76d917d02..84bf0cab1 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -2,6 +2,8 @@ # Native Asyncio Cursors +See {ref}`cursor-workers` for query and S3 worker settings and the 4.0.0 migration. + PyAthena provides native asyncio cursor implementations under `pyathena.aio`. These cursors use `asyncio.sleep` for polling and `asyncio.to_thread` for boto3 calls, keeping the event loop free. Concurrency comes from asyncio tasks rather than a diff --git a/docs/arrow.md b/docs/arrow.md index 5306f0902..6f01ecf8b 100644 --- a/docs/arrow.md +++ b/docs/arrow.md @@ -2,6 +2,8 @@ # Arrow +See {ref}`cursor-workers` for query and S3 worker settings and the 4.0.0 migration. + (arrow-cursor)= ## ArrowCursor diff --git a/docs/cursor.md b/docs/cursor.md index a3db13a2a..6f3c41fb7 100644 --- a/docs/cursor.md +++ b/docs/cursor.md @@ -1,5 +1,50 @@ # Cursor +(cursor-workers)= + +## Worker settings in 4.0.0 + +The thread-pool cursors use `max_workers` for their pool of tasks that wait for queries and collect results. +Query submission happens before the task is queued, so this setting does not cap the number of queries running in Athena. +The pandas, Arrow, and Polars cursors use `s3_max_workers` for the workers of each PyAthena S3 file reader. +It does not set a shared limit across queries or change the parsing library's CPU thread pool. + +| Cursor | `max_workers` | `s3_max_workers` default | +|--------|---------------|--------------------------| +| `PandasCursor`, `PolarsCursor`, `AioPandasCursor`, `AioPolarsCursor` | No query pool | `(cpu_count() or 1) * 5` | +| `AsyncPandasCursor`, `AsyncPolarsCursor` | Query task pool | `(cpu_count() or 1) * 5` | +| `ArrowCursor`, `AioArrowCursor` | No query pool | `None` (native PyArrow S3 filesystem) | +| `AsyncArrowCursor` | Query task pool | `None` (native PyArrow S3 filesystem) | + +Pass `s3_max_workers` to `execute()` to override the setting for one query. +The next call without an override uses the cursor's setting again. +For Arrow, a positive integer selects PyAthena's S3 filesystem through PyArrow's fsspec adapter for both CSV and UNLOAD results. +An explicit `None` selects the native PyArrow filesystem for that query. +On the PyAthena path, Arrow's `request_timeout` sets the boto3 read timeout; `connect_timeout` sets the connection timeout. +The connection's S3 configuration supplies defaults for unspecified timeouts. + +```python +from pyathena.pandas.async_cursor import AsyncPandasCursor + +cursor = connection.cursor(AsyncPandasCursor, max_workers=4, s3_max_workers=2) +query_id, future = cursor.execute("SELECT 1", s3_max_workers=3) +result = future.result() +``` + +The S3 setting applies when PyAthena's S3 filesystem reads the results. +User-provided pandas `filesystem` or `storage_options` can replace that filesystem and its settings. +Polars uses PyAthena's filesystem for CSV results; its native Parquet reader uses its own concurrency settings. + +### Migrate the former S3 argument + +This is a breaking API change in 4.0.0. +Replace `max_workers=N` with `s3_max_workers=N` in synchronous and native asyncio pandas/Polars cursor constructors, and in pandas/Polars `execute()` calls. +Keep `max_workers` in thread-pool cursor constructors to size the query task pool. +For `AsyncPolarsCursor`, specify both arguments to preserve the former behavior of controlling both pools with one value. +`AsyncPandasCursor` now also accepts an S3 worker default in its constructor. +The former S3 keyword raises `TypeError`; nonpositive S3 worker values raise `ValueError` before a query starts. +The `max_workers` names on result sets, S3 filesystems, and pandas upload utilities retain their existing meanings. + (default_cursor)= ## DefaultCursor diff --git a/docs/pandas.md b/docs/pandas.md index f4816ff49..5e43c33b5 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -2,6 +2,8 @@ # Pandas +See {ref}`cursor-workers` for query and S3 worker settings and the 4.0.0 migration. + (as-dataframe)= ## As DataFrame diff --git a/docs/polars.md b/docs/polars.md index f347e4214..8fedba7d0 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -2,6 +2,8 @@ # Polars +See {ref}`cursor-workers` for query and S3 worker settings and the 4.0.0 migration. + (polars-cursor)= ## PolarsCursor diff --git a/pyathena/aio/arrow/cursor.py b/pyathena/aio/arrow/cursor.py index c297dd0e1..f91d220e0 100644 --- a/pyathena/aio/arrow/cursor.py +++ b/pyathena/aio/arrow/cursor.py @@ -17,7 +17,7 @@ from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions -from pyathena.util import override +from pyathena.util import _validate_s3_max_workers, override if TYPE_CHECKING: import polars as pl @@ -55,6 +55,7 @@ def __init__( result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, connect_timeout: float | None = None, request_timeout: float | None = None, + s3_max_workers: int | None = None, **kwargs, ) -> None: """Initialize an AioArrowCursor. @@ -72,13 +73,17 @@ def __init__( unload: Whether to wrap queries in ``UNLOAD`` and read the Parquet output. result_reuse_enable: Whether to enable Athena query result reuse. result_reuse_minutes: Maximum age of a reused query result in minutes. - connect_timeout: Connection timeout in seconds of the pyarrow S3 filesystem - that reads the results. If None, the pyarrow default is used. - request_timeout: Request timeout in seconds of the pyarrow S3 filesystem - that reads the results. If None, the pyarrow default is used. + connect_timeout: S3 connection timeout in seconds. If None, use the + selected filesystem's default: PyArrow's default on the native path, + or the connection's S3 configuration on the PyAthena path. + request_timeout: S3 request timeout on the native PyArrow path, or boto3 + read timeout on the PyAthena path. If None, use that path's default. + s3_max_workers: S3 read workers per file. If None, use the native + pyarrow filesystem; otherwise use PyAthena's S3 filesystem. **kwargs: Other cursor arguments, such as ``connection`` and ``arraysize``, passed to the parent ``__init__``. """ + _validate_s3_max_workers(s3_max_workers, kwargs, allow_none=True) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -92,6 +97,7 @@ def __init__( result_reuse_minutes=result_reuse_minutes, **kwargs, ) + self._s3_max_workers = s3_max_workers self._unload = unload self._connect_timeout = connect_timeout self._request_timeout = request_timeout @@ -144,7 +150,8 @@ async def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Additional execution parameters. ``s3_max_workers`` + overrides the cursor's S3 read setting for this query. ``block_size`` sets the read block size for this query, and ``connect_timeout`` and ``request_timeout`` override the cursor's values. @@ -152,6 +159,9 @@ async def execute( Self reference for method chaining. """ self._reset_state() + _validate_s3_max_workers( + kwargs.get("s3_max_workers", self._s3_max_workers), kwargs, allow_none=True + ) options = ExecuteOptions.resolve( options, work_group=work_group, @@ -187,6 +197,7 @@ async def execute( unload_location=unload_location, connect_timeout=kwargs.pop("connect_timeout", self._connect_timeout), request_timeout=kwargs.pop("request_timeout", self._request_timeout), + s3_max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), result_set_type_hints=options.result_set_type_hints, **kwargs, ) diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index 4000d4baa..c052fc882 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -22,7 +22,7 @@ DefaultPandasUnloadTypeConverter, ) from pyathena.pandas.result_set import AthenaPandasResultSet, PandasDataFrameIterator -from pyathena.util import override +from pyathena.util import _validate_s3_max_workers, override if TYPE_CHECKING: from pandas import DataFrame @@ -60,7 +60,7 @@ def __init__( chunksize: int | None = None, block_size: int | None = None, cache_type: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + s3_max_workers: int = (cpu_count() or 1) * 5, result_reuse_enable: bool = False, result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, auto_optimize_chunksize: bool = False, @@ -84,7 +84,7 @@ def __init__( it takes precedence over ``auto_optimize_chunksize``. block_size: Default block size of the S3 filesystem that reads the results. cache_type: Default cache type of the S3 filesystem that reads the results. - max_workers: Maximum number of workers of the S3 filesystem. + s3_max_workers: Maximum number of workers of the S3 filesystem. result_reuse_enable: Whether to enable Athena query result reuse. result_reuse_minutes: Maximum age of a reused query result in minutes. auto_optimize_chunksize: Whether to choose a chunk size from the size of the @@ -92,6 +92,7 @@ def __init__( **kwargs: Other cursor arguments, such as ``connection`` and ``arraysize``, passed to the parent ``__init__``. """ + _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -110,7 +111,7 @@ def __init__( self._chunksize = chunksize self._block_size = block_size self._cache_type = cache_type - self._max_workers = max_workers + self._s3_max_workers = s3_max_workers self._auto_optimize_chunksize = auto_optimize_chunksize self._result_set: AthenaPandasResultSet | None = None @@ -168,7 +169,7 @@ async def execute( :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. **kwargs: Additional pandas read_csv/read_parquet parameters. - ``engine``, ``chunksize``, ``block_size``, ``cache_type``, ``max_workers``, + ``engine``, ``chunksize``, ``block_size``, ``cache_type``, ``s3_max_workers``, and ``auto_optimize_chunksize`` override the cursor's values for this query. ``storage_options`` and, for UNLOAD results, ``filesystem`` replace PyAthena's S3 filesystem (see @@ -178,6 +179,7 @@ async def execute( Self reference for method chaining. """ self._reset_state() + _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, @@ -218,7 +220,7 @@ async def execute( chunksize=kwargs.pop("chunksize", self._chunksize), block_size=kwargs.pop("block_size", self._block_size), cache_type=kwargs.pop("cache_type", self._cache_type), - max_workers=kwargs.pop("max_workers", self._max_workers), + max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), auto_optimize_chunksize=kwargs.pop( "auto_optimize_chunksize", self._auto_optimize_chunksize ), diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index ff9c7a60c..7f32cdf36 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -18,7 +18,7 @@ DefaultPolarsUnloadTypeConverter, ) from pyathena.polars.result_set import AthenaPolarsResultSet -from pyathena.util import override +from pyathena.util import _validate_s3_max_workers, override if TYPE_CHECKING: import polars as pl @@ -57,7 +57,7 @@ def __init__( result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, block_size: int | None = None, cache_type: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + s3_max_workers: int = (cpu_count() or 1) * 5, chunksize: int | None = None, **kwargs, ) -> None: @@ -78,12 +78,13 @@ def __init__( result_reuse_minutes: Maximum age of a reused query result in minutes. block_size: Default block size of the S3 filesystem that reads the results. cache_type: Default cache type of the S3 filesystem that reads the results. - max_workers: Maximum number of workers of the S3 filesystem. + s3_max_workers: Maximum number of workers of the S3 filesystem. chunksize: Number of rows per chunk. If set, result files in S3 are read lazily in chunks of this size. **kwargs: Other cursor arguments, such as ``connection`` and ``arraysize``, passed to the parent ``__init__``. """ + _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -100,7 +101,7 @@ def __init__( self._unload = unload self._block_size = block_size self._cache_type = cache_type - self._max_workers = max_workers + self._s3_max_workers = s3_max_workers self._chunksize = chunksize self._result_set: AthenaPolarsResultSet | None = None @@ -152,7 +153,7 @@ async def execute( :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. **kwargs: Additional execution parameters passed to Polars read functions. - ``block_size``, ``cache_type``, ``max_workers``, and ``chunksize`` + ``block_size``, ``cache_type``, ``s3_max_workers``, and ``chunksize`` override the cursor's values for this query. Read function arguments replace the ones the result set chooses, such as ``separator``, ``has_header``, ``schema_overrides``, and ``storage_options`` @@ -162,6 +163,7 @@ async def execute( Self reference for method chaining. """ self._reset_state() + _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, @@ -197,7 +199,7 @@ async def execute( unload_location=unload_location, block_size=kwargs.pop("block_size", self._block_size), cache_type=kwargs.pop("cache_type", self._cache_type), - max_workers=kwargs.pop("max_workers", self._max_workers), + max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), chunksize=kwargs.pop("chunksize", self._chunksize), result_set_type_hints=options.result_set_type_hints, **kwargs, diff --git a/pyathena/arrow/async_cursor.py b/pyathena/arrow/async_cursor.py index 4be5972fd..4adb1884e 100644 --- a/pyathena/arrow/async_cursor.py +++ b/pyathena/arrow/async_cursor.py @@ -17,7 +17,7 @@ from pyathena.common import CursorIterator from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions -from pyathena.util import override +from pyathena.util import _validate_s3_max_workers, override _logger = logging.getLogger(__name__) @@ -80,6 +80,7 @@ def __init__( result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, connect_timeout: float | None = None, request_timeout: float | None = None, + s3_max_workers: int | None = None, **kwargs, ) -> None: """Initialize an AsyncArrowCursor. @@ -93,17 +94,19 @@ def __init__( encryption_option: S3 encryption option (SSE_S3, SSE_KMS, CSE_KMS). kms_key: KMS key ARN for encryption. kill_on_interrupt: Cancel running query on keyboard interrupt. - max_workers: Maximum number of workers for concurrent execution. + max_workers: Size of the cursor thread pool for waiting and collecting results. arraysize: Number of rows to fetch per batch. unload: Enable UNLOAD for high-performance Parquet output. result_reuse_enable: Enable Athena query result reuse. result_reuse_minutes: Minutes to reuse cached results. - connect_timeout: Socket connection timeout in seconds for S3 operations. - Defaults to AWS SDK default (typically 1 second) if not specified. - request_timeout: Request timeout in seconds for S3 operations. - Defaults to AWS SDK default (typically 3 seconds) if not specified. - Increase this value if you experience timeout errors when using - role assumption with STS or have high latency to S3. + connect_timeout: S3 connection timeout in seconds. If None, use the + selected filesystem's default: PyArrow's default on the native path, + or the connection's S3 configuration on the PyAthena path. + request_timeout: S3 request timeout in seconds on the native PyArrow path, + or boto3 read timeout on the PyAthena path. If None, use the selected + filesystem's default. + s3_max_workers: S3 read workers per file. If None, use the native + pyarrow filesystem; otherwise use PyAthena's S3 filesystem. **kwargs: Additional connection parameters. Example: @@ -114,6 +117,7 @@ def __init__( ... request_timeout=30.0 ... ) """ + _validate_s3_max_workers(s3_max_workers, kwargs, allow_none=True) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -129,6 +133,7 @@ def __init__( result_reuse_minutes=result_reuse_minutes, **kwargs, ) + self._s3_max_workers = s3_max_workers self._unload = unload self._connect_timeout = connect_timeout self._request_timeout = request_timeout @@ -174,6 +179,7 @@ def _collect_result_set( unload_location=unload_location, connect_timeout=kwargs.pop("connect_timeout", self._connect_timeout), request_timeout=kwargs.pop("request_timeout", self._request_timeout), + s3_max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), result_set_type_hints=result_set_type_hints, **kwargs, ) @@ -212,13 +218,17 @@ def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Additional execution parameters. ``s3_max_workers`` + overrides the cursor's S3 read setting for this query. ``block_size`` sets the read block size for this query, and ``connect_timeout`` and ``request_timeout`` override the cursor's values. Returns: Tuple of (query_id, future) where future resolves to AthenaArrowResultSet. """ + _validate_s3_max_workers( + kwargs.get("s3_max_workers", self._s3_max_workers), kwargs, allow_none=True + ) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/arrow/cursor.py b/pyathena/arrow/cursor.py index e39c58237..4fe9798dc 100644 --- a/pyathena/arrow/cursor.py +++ b/pyathena/arrow/cursor.py @@ -16,7 +16,7 @@ from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions from pyathena.result_set import WithFetch -from pyathena.util import override +from pyathena.util import _validate_s3_max_workers, override if TYPE_CHECKING: import polars as pl @@ -70,6 +70,7 @@ def __init__( result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, connect_timeout: float | None = None, request_timeout: float | None = None, + s3_max_workers: int | None = None, **kwargs, ) -> None: """Initialize an ArrowCursor. @@ -86,12 +87,14 @@ def __init__( unload: Enable UNLOAD for high-performance Parquet output. result_reuse_enable: Enable Athena query result reuse. result_reuse_minutes: Minutes to reuse cached results. - connect_timeout: Socket connection timeout in seconds for S3 operations. - Defaults to AWS SDK default (typically 1 second) if not specified. - request_timeout: Request timeout in seconds for S3 operations. - Defaults to AWS SDK default (typically 3 seconds) if not specified. - Increase this value if you experience timeout errors when using - role assumption with STS or have high latency to S3. + connect_timeout: S3 connection timeout in seconds. If None, use the + selected filesystem's default: PyArrow's default on the native path, + or the connection's S3 configuration on the PyAthena path. + request_timeout: S3 request timeout in seconds on the native PyArrow path, + or boto3 read timeout on the PyAthena path. If None, use the selected + filesystem's default. + s3_max_workers: S3 read workers per file. If None, use the native + pyarrow filesystem; otherwise use PyAthena's S3 filesystem. **kwargs: Additional connection parameters. Example: @@ -102,6 +105,7 @@ def __init__( ... request_timeout=30 ... ) """ + _validate_s3_max_workers(s3_max_workers, kwargs, allow_none=True) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -115,6 +119,7 @@ def __init__( result_reuse_minutes=result_reuse_minutes, **kwargs, ) + self._s3_max_workers = s3_max_workers self._unload = unload self._connect_timeout = connect_timeout self._request_timeout = request_timeout @@ -170,7 +175,8 @@ def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Additional execution parameters. ``s3_max_workers`` + overrides the cursor's S3 read setting for this query. ``block_size`` sets the read block size for this query, and ``connect_timeout`` and ``request_timeout`` override the cursor's values. @@ -182,6 +188,9 @@ def execute( >>> table = cursor.as_arrow() # Returns Apache Arrow Table """ self._reset_state() + _validate_s3_max_workers( + kwargs.get("s3_max_workers", self._s3_max_workers), kwargs, allow_none=True + ) options = ExecuteOptions.resolve( options, work_group=work_group, @@ -215,6 +224,7 @@ def execute( unload_location=unload_location, connect_timeout=kwargs.pop("connect_timeout", self._connect_timeout), request_timeout=kwargs.pop("request_timeout", self._request_timeout), + s3_max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), result_set_type_hints=options.result_set_type_hints, **kwargs, ) diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 2b1b9cbb9..8fe01030f 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -4,19 +4,23 @@ import logging from collections.abc import Callable +from contextlib import ExitStack from typing import ( TYPE_CHECKING, Any, ClassVar, ) +from botocore.config import Config + from pyathena import OperationalError from pyathena.arrow.converter import _to_timestamp from pyathena.arrow.util import to_column_info from pyathena.converter import Converter, _to_default +from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet -from pyathena.util import RetryConfig, override, parse_output_location +from pyathena.util import RetryConfig, _validate_s3_max_workers, override, parse_output_location if TYPE_CHECKING: import polars as pl @@ -94,6 +98,8 @@ def __init__( connect_timeout: float | None = None, request_timeout: float | None = None, result_set_type_hints: dict[str | int, str] | None = None, + *, + s3_max_workers: int | None = None, **kwargs, ) -> None: """Initialize the result set and load the query results into an Arrow Table. @@ -111,16 +117,21 @@ def __init__( instead of the CSV results. unload_location: The S3 location of the ``UNLOAD`` output. If None, it is derived from the first file in the data manifest. - connect_timeout: The connect timeout in seconds for the pyarrow S3 filesystem. - request_timeout: The request timeout in seconds for the pyarrow S3 filesystem. + connect_timeout: Connection timeout in seconds for the selected S3 filesystem. + request_timeout: Native PyArrow request timeout, or boto3 read timeout + on the PyAthena path, in seconds. result_set_type_hints: Athena type signatures for complex-type columns, keyed by column name (case-insensitive) or zero-based column index. + s3_max_workers: S3 read workers per file. None uses the native pyarrow + filesystem. A positive integer uses PyAthena's S3 filesystem; + request_timeout then sets the boto3 read timeout. **kwargs: Additional keyword arguments, stored but not used. Raises: ProgrammingError: If ``query_execution`` is not given. OperationalError: If reading the query results fails. """ + _validate_s3_max_workers(s3_max_workers, kwargs, allow_none=True) super().__init__( connection=connection, converter=converter, @@ -137,21 +148,24 @@ def __init__( self._connect_timeout = connect_timeout self._request_timeout = request_timeout self._kwargs = kwargs - self._fs = self._create_s3_file_system() - if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location: - self._table = self._as_arrow() - elif self.state == AthenaQueryExecution.STATE_SUCCEEDED: - # Without a result file, as with managed query result storage, the rows from - # GetQueryResults are read as a CSV result file. - self._table = self._read_csv() - else: - import pyarrow as pa + self._s3_max_workers = s3_max_workers + self._s3_resources = ExitStack() + with self._s3_resources: + self._fs = self._create_s3_file_system() + if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location: + self._table = self._as_arrow() + elif self.state == AthenaQueryExecution.STATE_SUCCEEDED: + # Without a result file, as with managed query result storage, the rows from + # GetQueryResults are read as a CSV result file. + self._table = self._read_csv() + else: + import pyarrow as pa - self._table = pa.Table.from_pydict({}) + self._table = pa.Table.from_pydict({}) self._batches = iter(self._table.to_batches(arraysize)) def _create_s3_file_system(self): - """Create a pyarrow ``S3FileSystem`` from the connection settings. + """Create the selected S3 filesystem from the connection settings. Returns: The pyarrow S3 filesystem for reading the query results. @@ -160,6 +174,29 @@ def _create_s3_file_system(self): connection = self.connection + if self._s3_max_workers is not None: + client = None + if self._connect_timeout is not None or self._request_timeout is not None: + overrides = {} + if self._connect_timeout is not None: + overrides["connect_timeout"] = self._connect_timeout + if self._request_timeout is not None: + overrides["read_timeout"] = self._request_timeout + client = connection.session.client( + "s3", + region_name=connection.region_name, + config=connection.s3_config.merge(Config(**overrides)), + **connection._s3_client_kwargs, + ) + self._s3_resources.callback(client.close) + filesystem = S3FileSystem( + connection=connection, + max_workers=self._s3_max_workers, + s3_client=client, + skip_instance_cache=True, + ) + return fs.PyFileSystem(fs.FSSpecHandler(filesystem)) + # Build timeout parameters dict timeout_kwargs = {} if self._connect_timeout is not None: @@ -368,20 +405,23 @@ def _read_csv(self) -> Table: return pa.Table.from_pydict({}) try: - table = csv.read_csv( - self._fs.open_input_stream(location) if data is None else pa.BufferReader(data), - read_options=read_opts, - parse_options=parse_opts, - convert_options=csv.ConvertOptions( - strings_can_be_null=bool(binary_columns), - quoted_strings_can_be_null=False, - timestamp_parsers=self.timestamp_parsers, - column_types={ - **column_types, - **{column_names[i]: pa.string() for i in timestamp_types}, - }, - ), - ) + with ( + self._fs.open_input_stream(location) if data is None else pa.BufferReader(data) + ) as stream: + table = csv.read_csv( + stream, + read_options=read_opts, + parse_options=parse_opts, + convert_options=csv.ConvertOptions( + strings_can_be_null=bool(binary_columns), + quoted_strings_can_be_null=False, + timestamp_parsers=self.timestamp_parsers, + column_types={ + **column_types, + **{column_names[i]: pa.string() for i in timestamp_types}, + }, + ), + ) for index, type_ in timestamp_types.items(): table = table.set_column( index, diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index 825662d45..55ded4a63 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -84,7 +84,7 @@ def __init__( kill_on_interrupt: Cancel a query whose start in ``execute()`` is interrupted by ``KeyboardInterrupt``. Waiting runs on worker threads, which do not receive the interrupt. - max_workers: Maximum number of threads in the cursor's thread pool. + max_workers: Size of the cursor thread pool for waiting and collecting results. arraysize: Default number of rows per ``fetchmany()`` call of the result sets the cursor creates. result_reuse_enable: Enable Athena query result reuse. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index d3cf4537c..4ab7b2061 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -174,6 +174,7 @@ def __init__( allow_bucket_deletion: bool = False, version_aware: bool = False, *args, + s3_client: BaseClient | None = None, **kwargs, ) -> None: """Create a filesystem for Amazon S3. @@ -197,6 +198,9 @@ def __init__( allow_bucket_deletion: Whether ``rmdir`` may delete a bucket. version_aware: Whether reads pin the object version observed at open time. + s3_client: Optional S3 client overriding the connection or client + factory. The caller owns this client; the filesystem uses + the connection's retry policy when a connection is provided. *args: Passed to ``fsspec.AbstractFileSystem``. **kwargs: Passed to ``fsspec.AbstractFileSystem``; without a ``connection``, also s3fs-compatible client arguments. @@ -204,12 +208,13 @@ def __init__( the operations that accept ``RequestPayer``. """ super().__init__(*args, **kwargs) - if connection: + if s3_client is not None: + client = s3_client + elif connection: client = connection.s3_client - retry_config = connection.retry_config else: client = self._get_client_compatible_with_s3fs(**kwargs) - retry_config = RetryConfig() + retry_config = connection.retry_config if connection else RetryConfig() self.default_block_size = ( default_block_size if default_block_size else self.DEFAULT_BLOCK_SIZE ) diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index 01c968052..d49906529 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -18,7 +18,7 @@ DefaultPandasUnloadTypeConverter, ) from pyathena.pandas.result_set import AthenaPandasResultSet -from pyathena.util import override +from pyathena.util import _validate_s3_max_workers, override _logger = logging.getLogger(__name__) @@ -83,6 +83,7 @@ def __init__( block_size: int | None = None, cache_type: str | None = None, auto_optimize_chunksize: bool = False, + s3_max_workers: int = (cpu_count() or 1) * 5, **kwargs, ) -> None: """Initialize an AsyncPandasCursor. @@ -98,7 +99,7 @@ def __init__( kill_on_interrupt: Cancel a query whose start in ``execute()`` is interrupted by ``KeyboardInterrupt``. Waiting runs on worker threads, which do not receive the interrupt. - max_workers: Maximum number of threads that run queries concurrently. + max_workers: Size of the cursor thread pool for waiting and collecting results. arraysize: Number of rows to fetch per batch. Must be a positive integer. unload: Whether to wrap queries in ``UNLOAD`` and read the Parquet output. engine: Parsing engine (``auto``, ``c``, ``python``, or ``pyarrow``). @@ -110,9 +111,11 @@ def __init__( cache_type: Default cache type of the S3 filesystem that reads the results. auto_optimize_chunksize: Whether to choose a chunk size from the size of the CSV result file when ``chunksize`` is None. + s3_max_workers: Maximum worker threads per S3 file reader. **kwargs: Other cursor arguments, such as ``connection`` and ``converter``, passed to ``AsyncCursor.__init__``. """ + _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -128,6 +131,7 @@ def __init__( result_reuse_minutes=result_reuse_minutes, **kwargs, ) + self._s3_max_workers = s3_max_workers self._unload = unload self._engine = engine self._chunksize = chunksize @@ -184,6 +188,7 @@ def _collect_result_set( chunksize=kwargs.pop("chunksize", self._chunksize), block_size=kwargs.pop("block_size", self._block_size), cache_type=kwargs.pop("cache_type", self._cache_type), + max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), auto_optimize_chunksize=kwargs.pop( "auto_optimize_chunksize", self._auto_optimize_chunksize ), @@ -237,11 +242,12 @@ def execute( ``storage_options`` and, for UNLOAD results, ``filesystem`` replace PyAthena's S3 filesystem (see :class:`~pyathena.pandas.result_set.AthenaPandasResultSet`). - ``max_workers`` sets the number of S3 read workers for this query. + ``s3_max_workers`` sets the number of S3 read workers for this query. Returns: Tuple of (query_id, future) where future resolves to AthenaPandasResultSet. """ + _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 1ed437935..0d542bf6b 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -21,7 +21,7 @@ ) from pyathena.pandas.result_set import AthenaPandasResultSet, PandasDataFrameIterator from pyathena.result_set import WithFetch -from pyathena.util import override +from pyathena.util import _validate_s3_max_workers, override if TYPE_CHECKING: from pandas import DataFrame @@ -81,7 +81,7 @@ def __init__( chunksize: int | None = None, block_size: int | None = None, cache_type: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + s3_max_workers: int = (cpu_count() or 1) * 5, result_reuse_enable: bool = False, result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, auto_optimize_chunksize: bool = False, @@ -104,7 +104,7 @@ def __init__( If specified, takes precedence over auto_optimize_chunksize. block_size: S3 read block size. cache_type: S3 caching strategy. - max_workers: Maximum worker threads for parallel processing. + s3_max_workers: Maximum worker threads per S3 file reader. result_reuse_enable: Enable query result reuse. result_reuse_minutes: Result reuse duration in minutes. auto_optimize_chunksize: Enable automatic chunksize determination for @@ -115,6 +115,7 @@ def __init__( ``converter``, ``formatter``, and ``retry_config``. Pass pandas ``read_csv``/``read_parquet`` options to ``execute()`` instead. """ + _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -133,7 +134,7 @@ def __init__( self._chunksize = chunksize self._block_size = block_size self._cache_type = cache_type - self._max_workers = max_workers + self._s3_max_workers = s3_max_workers self._auto_optimize_chunksize = auto_optimize_chunksize @staticmethod @@ -194,7 +195,7 @@ def execute( :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. **kwargs: Additional pandas read_csv/read_parquet parameters. - ``engine``, ``chunksize``, ``block_size``, ``cache_type``, ``max_workers``, + ``engine``, ``chunksize``, ``block_size``, ``cache_type``, ``s3_max_workers``, and ``auto_optimize_chunksize`` override the cursor's values for this query. ``storage_options`` and, for UNLOAD results, ``filesystem`` replace PyAthena's S3 filesystem (see @@ -209,6 +210,7 @@ def execute( >>> df = cursor.as_pandas() # Returns pandas DataFrame """ self._reset_state() + _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, @@ -247,7 +249,7 @@ def execute( chunksize=kwargs.pop("chunksize", self._chunksize), block_size=kwargs.pop("block_size", self._block_size), cache_type=kwargs.pop("cache_type", self._cache_type), - max_workers=kwargs.pop("max_workers", self._max_workers), + max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), auto_optimize_chunksize=kwargs.pop( "auto_optimize_chunksize", self._auto_optimize_chunksize ), diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index f836b2179..85bf660ea 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -17,7 +17,7 @@ DefaultPolarsUnloadTypeConverter, ) from pyathena.polars.result_set import AthenaPolarsResultSet -from pyathena.util import override +from pyathena.util import _validate_s3_max_workers, override _logger = logging.getLogger(__name__) @@ -78,6 +78,7 @@ def __init__( block_size: int | None = None, cache_type: str | None = None, chunksize: int | None = None, + s3_max_workers: int = (cpu_count() or 1) * 5, **kwargs, ) -> None: """Initialize an AsyncPolarsCursor. @@ -91,7 +92,7 @@ def __init__( encryption_option: S3 encryption option (SSE_S3, SSE_KMS, CSE_KMS). kms_key: KMS key ARN for encryption. kill_on_interrupt: Cancel running query on keyboard interrupt. - max_workers: Maximum number of workers for concurrent execution. + max_workers: Size of the cursor thread pool for waiting and collecting results. arraysize: Number of rows to fetch per batch. unload: Enable UNLOAD for high-performance Parquet output. result_reuse_enable: Enable Athena query result reuse. @@ -101,6 +102,7 @@ def __init__( chunksize: Number of rows per chunk for memory-efficient processing. If specified, data is loaded lazily in chunks for all data access methods including fetchone(), fetchmany(), and iter_chunks(). + s3_max_workers: Maximum worker threads per S3 file reader. **kwargs: Additional connection parameters. Example: @@ -108,6 +110,7 @@ def __init__( >>> # With chunked processing >>> cursor = connection.cursor(AsyncPolarsCursor, chunksize=50000) """ + _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -123,6 +126,7 @@ def __init__( result_reuse_minutes=result_reuse_minutes, **kwargs, ) + self._s3_max_workers = s3_max_workers self._unload = unload self._block_size = block_size self._cache_type = cache_type @@ -186,7 +190,7 @@ def _collect_result_set( unload_location=unload_location, block_size=kwargs.pop("block_size", self._block_size), cache_type=kwargs.pop("cache_type", self._cache_type), - max_workers=kwargs.pop("max_workers", self._max_workers), + max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), chunksize=kwargs.pop("chunksize", self._chunksize), result_set_type_hints=result_set_type_hints, **kwargs, @@ -231,7 +235,7 @@ def execute( :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. **kwargs: Additional execution parameters passed to Polars read functions. - ``block_size``, ``cache_type``, ``max_workers``, and ``chunksize`` + ``block_size``, ``cache_type``, ``s3_max_workers``, and ``chunksize`` override the cursor's values for this query. Read function arguments replace the ones the result set chooses, such as ``separator``, ``has_header``, ``schema_overrides``, and ``storage_options`` @@ -245,6 +249,7 @@ def execute( >>> result_set = future.result() >>> df = result_set.as_polars() # Returns Polars DataFrame """ + _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index a7e31a9f2..5e5811636 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -21,7 +21,7 @@ ) from pyathena.polars.result_set import AthenaPolarsResultSet from pyathena.result_set import WithFetch -from pyathena.util import override +from pyathena.util import _validate_s3_max_workers, override if TYPE_CHECKING: import polars as pl @@ -81,7 +81,7 @@ def __init__( result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, block_size: int | None = None, cache_type: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + s3_max_workers: int = (cpu_count() or 1) * 5, chunksize: int | None = None, **kwargs, ) -> None: @@ -101,7 +101,7 @@ def __init__( result_reuse_minutes: Minutes to reuse cached results. block_size: S3 read block size. cache_type: S3 caching strategy. - max_workers: Maximum worker threads for parallel S3 operations. + s3_max_workers: Maximum worker threads for parallel S3 operations. chunksize: Number of rows per chunk for memory-efficient processing. If specified, data is loaded lazily in chunks for all data access methods including fetchone(), fetchmany(), and iter_chunks(). @@ -112,6 +112,7 @@ def __init__( >>> # With chunked processing >>> cursor = connection.cursor(PolarsCursor, chunksize=50000) """ + _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -128,7 +129,7 @@ def __init__( self._unload = unload self._block_size = block_size self._cache_type = cache_type - self._max_workers = max_workers + self._s3_max_workers = s3_max_workers self._chunksize = chunksize @staticmethod @@ -190,7 +191,7 @@ def execute( :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. **kwargs: Additional execution parameters passed to Polars read functions. - ``block_size``, ``cache_type``, ``max_workers``, and ``chunksize`` + ``block_size``, ``cache_type``, ``s3_max_workers``, and ``chunksize`` override the cursor's values for this query. Read function arguments replace the ones the result set chooses, such as ``separator``, ``has_header``, ``schema_overrides``, and ``storage_options`` @@ -204,6 +205,7 @@ def execute( >>> df = cursor.as_polars() # Returns Polars DataFrame """ self._reset_state() + _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, @@ -237,7 +239,7 @@ def execute( unload_location=unload_location, block_size=kwargs.pop("block_size", self._block_size), cache_type=kwargs.pop("cache_type", self._cache_type), - max_workers=kwargs.pop("max_workers", self._max_workers), + max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), chunksize=kwargs.pop("chunksize", self._chunksize), result_set_type_hints=options.result_set_type_hints, **kwargs, diff --git a/pyathena/util.py b/pyathena/util.py index a88dc6bf5..05e5b4f61 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -54,6 +54,30 @@ def override(func): ) +def _validate_s3_max_workers( + value: int | None, kwargs: dict[str, Any], *, allow_none: bool = False +) -> None: + """Validate cursor S3 workers and reject the former S3 keyword. + + Args: + value: The S3 worker setting to validate. + kwargs: Cursor constructor or execution keyword arguments. + allow_none: Whether None selects the native Arrow filesystem. + + Raises: + TypeError: If max_workers is used as a read option, or value is not an integer. + ValueError: If value is not positive. + """ + if "max_workers" in kwargs: + raise TypeError("Use s3_max_workers for S3 reads; max_workers configures async cursors.") + if value is None and allow_none: + return + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError("s3_max_workers must be a positive integer") + if value <= 0: + raise ValueError("s3_max_workers must be greater than 0") + + def parse_output_location(output_location: str) -> tuple[str, str]: """Parse an S3 output location URL into bucket and key components. diff --git a/tests/pyathena/aio/arrow/test_cursor.py b/tests/pyathena/aio/arrow/test_cursor.py index fa4e14e11..67c611835 100644 --- a/tests/pyathena/aio/arrow/test_cursor.py +++ b/tests/pyathena/aio/arrow/test_cursor.py @@ -19,6 +19,20 @@ class TestAioArrowCursor: + @pytest.mark.parametrize( + "aio_arrow_cursor", + [ + {"cursor_kwargs": {"s3_max_workers": 2, "unload": False}}, + {"cursor_kwargs": {"s3_max_workers": 2, "unload": True}}, + ], + indirect=True, + ) + async def test_s3_workers_read_results(self, aio_arrow_cursor): + """CSV and Parquet results use the configured PyAthena S3 reader.""" + await aio_arrow_cursor.execute("SELECT 1 AS value", s3_max_workers=1) + assert await aio_arrow_cursor.fetchall() == [(1,)] + assert aio_arrow_cursor.result_set._fs.handler.fs.max_workers == 1 + async def test_binary_null_vs_empty(self, aio_arrow_cursor): query = """SELECT * FROM (VALUES (1, CAST(NULL AS VARBINARY), 'null', CAST(NULL AS VARCHAR)), @@ -150,14 +164,19 @@ async def test_as_arrow_unload(self, aio_arrow_cursor): assert table.num_rows == 1 @pytest.mark.parametrize( - "execute_kwargs", [{}, {"connect_timeout": 3.0, "request_timeout": 4.0}] + "execute_kwargs", + [ + {}, + {"connect_timeout": 3.0, "request_timeout": 4.0, "s3_max_workers": 3}, + {"s3_max_workers": None}, + ], ) async def test_read_options(self, execute_kwargs): """The cursor's read options reach the result set, and execute() overrides them. No AWS calls; the query and its result set are mocked. """ - cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0} + cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0, "s3_max_workers": 2} query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) cursor = AioArrowCursor( connection=MagicMock(), @@ -172,6 +191,12 @@ async def test_read_options(self, execute_kwargs): patch("pyathena.aio.arrow.cursor.AthenaArrowResultSet") as result_set_class, ): await cursor.execute("SELECT 1", **execute_kwargs) - kwargs = result_set_class.call_args.kwargs + first_kwargs = result_set_class.call_args.kwargs.copy() + await cursor.execute("SELECT 1") + kwargs = first_kwargs expected = {**cursor_kwargs, **execute_kwargs} assert {key: kwargs[key] for key in expected} == expected + + defaults = cursor_kwargs.copy() + second_kwargs = result_set_class.call_args.kwargs + assert {key: second_kwargs[key] for key in defaults} == defaults diff --git a/tests/pyathena/aio/pandas/test_cursor.py b/tests/pyathena/aio/pandas/test_cursor.py index ecb123a34..eebfb8778 100644 --- a/tests/pyathena/aio/pandas/test_cursor.py +++ b/tests/pyathena/aio/pandas/test_cursor.py @@ -147,7 +147,7 @@ async def test_as_pandas_unload(self, aio_pandas_cursor): { "block_size": 2048, "cache_type": "none", - "max_workers": 3, + "s3_max_workers": 3, "auto_optimize_chunksize": False, }, ], @@ -160,7 +160,7 @@ async def test_read_options(self, execute_kwargs): cursor_kwargs = { "block_size": 1024, "cache_type": "bytes", - "max_workers": 2, + "s3_max_workers": 2, "auto_optimize_chunksize": True, } cursor = AioPandasCursor( @@ -177,6 +177,14 @@ async def test_read_options(self, execute_kwargs): patch("pyathena.aio.pandas.cursor.AthenaPandasResultSet") as result_set_class, ): await cursor.execute("SELECT 1", **execute_kwargs) - kwargs = result_set_class.call_args.kwargs + first_kwargs = result_set_class.call_args.kwargs.copy() + await cursor.execute("SELECT 1") + kwargs = first_kwargs expected = {**cursor_kwargs, **execute_kwargs} + expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected + + defaults = cursor_kwargs.copy() + defaults["max_workers"] = defaults.pop("s3_max_workers") + second_kwargs = result_set_class.call_args.kwargs + assert {key: second_kwargs[key] for key in defaults} == defaults diff --git a/tests/pyathena/aio/polars/test_cursor.py b/tests/pyathena/aio/polars/test_cursor.py index 051d078dc..049bd38bf 100644 --- a/tests/pyathena/aio/polars/test_cursor.py +++ b/tests/pyathena/aio/polars/test_cursor.py @@ -136,7 +136,7 @@ async def test_as_polars_unload(self, aio_polars_cursor): @pytest.mark.parametrize( "execute_kwargs", - [{}, {"block_size": 2048, "cache_type": "none", "max_workers": 3, "chunksize": 20}], + [{}, {"block_size": 2048, "cache_type": "none", "s3_max_workers": 3, "chunksize": 20}], ) async def test_read_options(self, execute_kwargs): """The cursor's read options reach the result set, and execute() overrides them. @@ -146,7 +146,7 @@ async def test_read_options(self, execute_kwargs): cursor_kwargs = { "block_size": 1024, "cache_type": "bytes", - "max_workers": 2, + "s3_max_workers": 2, "chunksize": 10, } cursor = AioPolarsCursor( @@ -163,6 +163,14 @@ async def test_read_options(self, execute_kwargs): patch("pyathena.aio.polars.cursor.AthenaPolarsResultSet") as result_set_class, ): await cursor.execute("SELECT 1", **execute_kwargs) - kwargs = result_set_class.call_args.kwargs + first_kwargs = result_set_class.call_args.kwargs.copy() + await cursor.execute("SELECT 1") + kwargs = first_kwargs expected = {**cursor_kwargs, **execute_kwargs} + expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected + + defaults = cursor_kwargs.copy() + defaults["max_workers"] = defaults.pop("s3_max_workers") + second_kwargs = result_set_class.call_args.kwargs + assert {key: second_kwargs[key] for key in defaults} == defaults diff --git a/tests/pyathena/arrow/test_async_cursor.py b/tests/pyathena/arrow/test_async_cursor.py index 30ae8f5ce..5d4d9cb09 100644 --- a/tests/pyathena/arrow/test_async_cursor.py +++ b/tests/pyathena/arrow/test_async_cursor.py @@ -18,6 +18,21 @@ class TestAsyncArrowCursor: + @pytest.mark.parametrize( + "async_arrow_cursor", + [ + {"cursor_kwargs": {"s3_max_workers": 2, "unload": False}}, + {"cursor_kwargs": {"s3_max_workers": 2, "unload": True}}, + ], + indirect=True, + ) + def test_s3_workers_read_results(self, async_arrow_cursor): + """CSV and Parquet results use the configured PyAthena S3 reader.""" + _, future = async_arrow_cursor.execute("SELECT 1 AS value", s3_max_workers=1) + result = future.result() + assert result.fetchall() == [(1,)] + assert result._fs.handler.fs.max_workers == 1 + def test_binary_null_vs_empty(self, async_arrow_cursor): query = """SELECT * FROM (VALUES (1, CAST(NULL AS VARBINARY), 'null', CAST(NULL AS VARCHAR)), @@ -327,14 +342,19 @@ def test_empty_result_unload(self, async_arrow_cursor): assert table.shape[1] == 0 @pytest.mark.parametrize( - "execute_kwargs", [{}, {"connect_timeout": 3.0, "request_timeout": 4.0}] + "execute_kwargs", + [ + {}, + {"connect_timeout": 3.0, "request_timeout": 4.0, "s3_max_workers": 3}, + {"s3_max_workers": None}, + ], ) def test_read_options(self, execute_kwargs): """The cursor's read options reach the result set, and execute() overrides them. No AWS calls; the query and its result set are mocked. """ - cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0} + cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0, "s3_max_workers": 2} query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) with ( AsyncArrowCursor( @@ -342,6 +362,7 @@ def test_read_options(self, execute_kwargs): converter=MagicMock(), formatter=MagicMock(), retry_config=RetryConfig(), + max_workers=1, **cursor_kwargs, ) as cursor, patch.object(AsyncArrowCursor, "_execute", return_value="query_id"), @@ -350,6 +371,14 @@ def test_read_options(self, execute_kwargs): ): _, future = cursor.execute("SELECT 1", **execute_kwargs) future.result() - kwargs = result_set_class.call_args.kwargs + first_kwargs = result_set_class.call_args.kwargs.copy() + _, future = cursor.execute("SELECT 1") + future.result() + kwargs = first_kwargs expected = {**cursor_kwargs, **execute_kwargs} assert {key: kwargs[key] for key in expected} == expected + + defaults = cursor_kwargs.copy() + second_kwargs = result_set_class.call_args.kwargs + assert {key: second_kwargs[key] for key in defaults} == defaults + assert cursor._executor._max_workers == 1 diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index 54352fa9c..49196e459 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -35,6 +35,20 @@ class TestArrowCursor: + @pytest.mark.parametrize( + "arrow_cursor", + [ + {"cursor_kwargs": {"s3_max_workers": 2, "unload": False}}, + {"cursor_kwargs": {"s3_max_workers": 2, "unload": True}}, + ], + indirect=True, + ) + def test_s3_workers_read_results(self, arrow_cursor): + """CSV and Parquet results use the configured PyAthena S3 reader.""" + arrow_cursor.execute("SELECT 1 AS value", s3_max_workers=1) + assert arrow_cursor.fetchall() == [(1,)] + assert arrow_cursor.result_set._fs.handler.fs.max_workers == 1 + def test_binary_null_vs_empty(self, arrow_cursor): # The text column has the name of the binary column, and keeps its own NULL. query = """SELECT id, value, label, text_value AS value FROM (VALUES @@ -1126,14 +1140,19 @@ def test_duplicate_column_names_with_newline(self, arrow_cursor): assert arrow_cursor.as_arrow().column_names == ["a\nb", "a\nb"] @pytest.mark.parametrize( - "execute_kwargs", [{}, {"connect_timeout": 3.0, "request_timeout": 4.0}] + "execute_kwargs", + [ + {}, + {"connect_timeout": 3.0, "request_timeout": 4.0, "s3_max_workers": 3}, + {"s3_max_workers": None}, + ], ) def test_read_options(self, execute_kwargs): """The cursor's read options reach the result set, and execute() overrides them. No AWS calls; the query and its result set are mocked. """ - cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0} + cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0, "s3_max_workers": 2} query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) cursor = ArrowCursor( connection=MagicMock(), @@ -1148,6 +1167,12 @@ def test_read_options(self, execute_kwargs): patch("pyathena.arrow.cursor.AthenaArrowResultSet") as result_set_class, ): cursor.execute("SELECT 1", **execute_kwargs) - kwargs = result_set_class.call_args.kwargs + first_kwargs = result_set_class.call_args.kwargs.copy() + cursor.execute("SELECT 1") + kwargs = first_kwargs expected = {**cursor_kwargs, **execute_kwargs} assert {key: kwargs[key] for key in expected} == expected + + defaults = cursor_kwargs.copy() + second_kwargs = result_set_class.call_args.kwargs + assert {key: second_kwargs[key] for key in defaults} == defaults diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index 90cf33c9a..cfadf3d2e 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -4,14 +4,24 @@ # See LICENSE or https://opensource.org/licenses/MIT. # # SPDX-License-Identifier: MIT -from datetime import datetime +from datetime import UTC, datetime +from io import BytesIO from unittest.mock import MagicMock, patch +import boto3 import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from botocore.config import Config +from botocore.exceptions import ClientError +from botocore.response import StreamingBody +from fsspec.implementations.memory import MemoryFileSystem from pyathena.arrow.converter import DefaultArrowTypeConverter from pyathena.arrow.result_set import AthenaArrowResultSet +from pyathena.filesystem.s3_executor import S3ThreadPoolExecutor from pyathena.model import AthenaQueryExecution +from pyathena.result_set import AthenaResultSet from pyathena.util import RetryConfig @@ -54,3 +64,140 @@ def test_read_csv_timestamps_with_the_same_name(self): ("a", datetime(2020, 1, 2, 3, 4, 5, 123456), datetime(2020, 1, 2, 3, 4, 5, 100000)), ("b", None, None), ] + + @pytest.mark.parametrize("unload", [False, True]) + def test_s3_workers_read_files(self, unload): + """Read real CSV/Parquet bytes through the S3 file and Arrow adapter offline.""" + if unload: + output = pa.BufferOutputStream() + pq.write_table(pa.table({"value": [1]}), output) + data = output.getvalue().to_pybytes() + key = "unload/part.parquet" + else: + data = b'"value"\n1\n' + key = "result.csv" + model_client = boto3.client( + "s3", region_name="us-east-1", aws_access_key_id="dummy", aws_secret_access_key="dummy" + ) + client = MagicMock() + client.meta.service_model = model_client.meta.service_model + client.meta.method_to_api_mapping = model_client.meta.method_to_api_mapping + model_client.close() + + def head_object(**kwargs): + if kwargs["Key"] != key: + raise ClientError( + {"Error": {"Code": "404"}, "ResponseMetadata": {"HTTPStatusCode": 404}}, + "HeadObject", + ) + return {"ContentLength": len(data), "LastModified": datetime(2026, 1, 1, tzinfo=UTC)} + + def get_object(**kwargs): + start, end = map(int, kwargs["Range"].removeprefix("bytes=").split("-")) + body = data[start : end + 1] + return {"Body": StreamingBody(BytesIO(body), len(body))} + + client.head_object.side_effect = head_object + client.get_object.side_effect = get_object + client.list_objects_v2.return_value = { + "Contents": [{"Key": key, "Size": len(data)}], + "IsTruncated": False, + } + connection = MagicMock(s3_client=client, retry_config=RetryConfig(attempt=1)) + execution = MagicMock( + state=AthenaQueryExecution.STATE_FAILED, + output_location=f"s3://bucket/{key}", + substatement_type=None, + ) + result = AthenaArrowResultSet( + connection=connection, + converter=DefaultArrowTypeConverter(), + query_execution=execution, + arraysize=1, + retry_config=RetryConfig(), + s3_max_workers=2, + unload=unload, + unload_location="s3://bucket/unload/", + ) + execution.state = AthenaQueryExecution.STATE_SUCCEEDED + result._metadata = ( + {"Name": "value", "Type": "integer", "Precision": 0, "Scale": 0, "Nullable": "UNKNOWN"}, + ) + with ( + patch.object(result, "_read_data_manifest", return_value=[f"s3://bucket/{key}"]), + patch( + "pyathena.filesystem.s3.S3ThreadPoolExecutor", wraps=S3ThreadPoolExecutor + ) as executor, + ): + table = result._as_arrow() + assert table.to_pylist() == [{"value": 1}] + assert executor.call_count > 0 + assert all(call.kwargs["max_workers"] == 2 for call in executor.call_args_list) + assert result._fs.handler.fs.max_workers == 2 + client.get_object.assert_called() + client.close.assert_not_called() + + @pytest.mark.parametrize("failure", [None, "filesystem", "read"]) + def test_s3_workers_owned_client_cleanup(self, failure): + """Dedicated timeout clients close after eager reads, including failures.""" + connection = MagicMock(region_name="us-east-1") + connection.s3_config = Config(max_pool_connections=17, read_timeout=60) + connection._s3_client_kwargs = {"use_ssl": True} + owned_client = connection.session.client.return_value + filesystem = MemoryFileSystem(skip_instance_cache=True) + execution = MagicMock( + state=AthenaQueryExecution.STATE_SUCCEEDED, output_location="s3://bucket/result.csv" + ) + with ( + patch.object(AthenaResultSet, "_pre_fetch"), + patch("pyathena.arrow.result_set.S3FileSystem", return_value=filesystem) as factory, + patch.object( + AthenaArrowResultSet, "_as_arrow", return_value=pa.table({"value": [1]}) + ) as read, + ): + if failure == "filesystem": + factory.side_effect = RuntimeError("filesystem failed") + elif failure == "read": + read.side_effect = RuntimeError("read failed") + kwargs = { + "connection": connection, + "converter": DefaultArrowTypeConverter(), + "query_execution": execution, + "arraysize": 1, + "retry_config": RetryConfig(), + "s3_max_workers": 2, + "connect_timeout": 1.5, + "request_timeout": 4.5, + } + if failure: + with pytest.raises(RuntimeError, match="failed"): + AthenaArrowResultSet(**kwargs) + else: + result = AthenaArrowResultSet(**kwargs) + assert result.as_arrow().to_pylist() == [{"value": 1}] + config = connection.session.client.call_args.kwargs["config"] + assert config.connect_timeout == 1.5 + assert config.read_timeout == 4.5 + assert config.max_pool_connections == 17 + assert connection.s3_config.read_timeout == 60 + assert connection.session.client.call_args.kwargs["use_ssl"] is True + assert factory.call_args.kwargs["s3_client"] is owned_client + owned_client.close.assert_called_once_with() + connection.s3_client.close.assert_not_called() + + def test_s3_workers_native_default(self): + """The default still builds the native Arrow filesystem, without a boto3 client.""" + connection = MagicMock(region_name="us-east-1", profile_name=None) + connection._kwargs = {"aws_access_key_id": "dummy", "aws_secret_access_key": "dummy"} + with patch("pyarrow.fs.S3FileSystem") as native: + AthenaArrowResultSet( + connection=connection, + converter=DefaultArrowTypeConverter(), + query_execution=MagicMock(state=AthenaQueryExecution.STATE_FAILED), + arraysize=1, + retry_config=RetryConfig(), + ) + native.assert_called_once_with( + access_key="dummy", secret_key="dummy", session_token=None, region="us-east-1" + ) + connection.session.client.assert_not_called() diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 410c6f697..9051fc5b2 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -115,6 +115,20 @@ def __getitem__(self, item): return BarrierDirCache() + @pytest.mark.parametrize("with_connection", [False, True]) + def test_s3_workers_client_override(self, with_connection): + """An injected client bypasses client creation and preserves connection retries.""" + connection = mock.MagicMock() if with_connection else None + client = mock.MagicMock() + with mock.patch.object(S3FileSystem, "_get_client_compatible_with_s3fs") as factory: + fs = S3FileSystem(connection=connection, s3_client=client, skip_instance_cache=True) + factory.assert_not_called() + assert fs.core.client is client + if connection is not None: + assert fs.core.retry_config is connection.retry_config + assert "s3_client" not in vars(connection) + client.close.assert_not_called() + def test_get_client_compatible_with_s3fs(self): # Only constructs a boto3 client; no AWS access. fs = S3FileSystem( diff --git a/tests/pyathena/pandas/test_async_cursor.py b/tests/pyathena/pandas/test_async_cursor.py index 590d9d8b8..fce96b16f 100644 --- a/tests/pyathena/pandas/test_async_cursor.py +++ b/tests/pyathena/pandas/test_async_cursor.py @@ -637,7 +637,7 @@ def test_null_decimal_value(self, async_pandas_cursor, parquet_engine): { "block_size": 2048, "cache_type": "none", - "max_workers": 3, + "s3_max_workers": 3, "auto_optimize_chunksize": False, }, ], @@ -647,7 +647,12 @@ def test_read_options(self, execute_kwargs): No AWS calls; the query and its result set are mocked. """ - cursor_kwargs = {"block_size": 1024, "cache_type": "bytes", "auto_optimize_chunksize": True} + cursor_kwargs = { + "block_size": 1024, + "cache_type": "bytes", + "auto_optimize_chunksize": True, + "s3_max_workers": 2, + } query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) with ( AsyncPandasCursor( @@ -655,6 +660,7 @@ def test_read_options(self, execute_kwargs): converter=MagicMock(), formatter=MagicMock(), retry_config=RetryConfig(), + max_workers=1, **cursor_kwargs, ) as cursor, patch.object(AsyncPandasCursor, "_execute", return_value="query_id"), @@ -663,10 +669,20 @@ def test_read_options(self, execute_kwargs): ): _, future = cursor.execute("SELECT 1", **execute_kwargs) future.result() - kwargs = result_set_class.call_args.kwargs + first_kwargs = result_set_class.call_args.kwargs.copy() + _, future = cursor.execute("SELECT 1") + future.result() + kwargs = first_kwargs expected = {**cursor_kwargs, **execute_kwargs} + expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected + defaults = cursor_kwargs.copy() + defaults["max_workers"] = defaults.pop("s3_max_workers") + second_kwargs = result_set_class.call_args.kwargs + assert {key: second_kwargs[key] for key in defaults} == defaults + assert cursor._executor._max_workers == 1 + @pytest.mark.parametrize( "async_pandas_cursor", [{"cursor_kwargs": {"auto_optimize_chunksize": True}}], diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index ad7bce11d..1d353dd49 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -1908,7 +1908,7 @@ def test_duplicate_column_names_read_options( { "block_size": 2048, "cache_type": "none", - "max_workers": 3, + "s3_max_workers": 3, "auto_optimize_chunksize": False, }, ], @@ -1921,7 +1921,7 @@ def test_read_options(self, execute_kwargs): cursor_kwargs = { "block_size": 1024, "cache_type": "bytes", - "max_workers": 2, + "s3_max_workers": 2, "auto_optimize_chunksize": True, } cursor = PandasCursor( @@ -1938,10 +1938,18 @@ def test_read_options(self, execute_kwargs): patch("pyathena.pandas.cursor.AthenaPandasResultSet") as result_set_class, ): cursor.execute("SELECT 1", **execute_kwargs) - kwargs = result_set_class.call_args.kwargs + first_kwargs = result_set_class.call_args.kwargs.copy() + cursor.execute("SELECT 1") + kwargs = first_kwargs expected = {**cursor_kwargs, **execute_kwargs} + expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected + defaults = cursor_kwargs.copy() + defaults["max_workers"] = defaults.pop("s3_max_workers") + second_kwargs = result_set_class.call_args.kwargs + assert {key: second_kwargs[key] for key in defaults} == defaults + @pytest.mark.parametrize( "pandas_cursor", [ diff --git a/tests/pyathena/polars/test_async_cursor.py b/tests/pyathena/polars/test_async_cursor.py index ac4a761d9..9400eadf9 100644 --- a/tests/pyathena/polars/test_async_cursor.py +++ b/tests/pyathena/polars/test_async_cursor.py @@ -328,7 +328,7 @@ def test_empty_result_unload(self, async_polars_cursor): @pytest.mark.parametrize( "execute_kwargs", - [{}, {"block_size": 2048, "cache_type": "none", "max_workers": 3, "chunksize": 20}], + [{}, {"block_size": 2048, "cache_type": "none", "s3_max_workers": 3, "chunksize": 20}], ) def test_read_options(self, execute_kwargs): """The cursor's read options reach the result set, and execute() overrides them. @@ -338,7 +338,7 @@ def test_read_options(self, execute_kwargs): cursor_kwargs = { "block_size": 1024, "cache_type": "bytes", - "max_workers": 2, + "s3_max_workers": 2, "chunksize": 10, } query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) @@ -348,6 +348,7 @@ def test_read_options(self, execute_kwargs): converter=MagicMock(), formatter=MagicMock(), retry_config=RetryConfig(), + max_workers=1, **cursor_kwargs, ) as cursor, patch.object(AsyncPolarsCursor, "_execute", return_value="query_id"), @@ -356,6 +357,16 @@ def test_read_options(self, execute_kwargs): ): _, future = cursor.execute("SELECT 1", **execute_kwargs) future.result() - kwargs = result_set_class.call_args.kwargs + first_kwargs = result_set_class.call_args.kwargs.copy() + _, future = cursor.execute("SELECT 1") + future.result() + kwargs = first_kwargs expected = {**cursor_kwargs, **execute_kwargs} + expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected + + defaults = cursor_kwargs.copy() + defaults["max_workers"] = defaults.pop("s3_max_workers") + second_kwargs = result_set_class.call_args.kwargs + assert {key: second_kwargs[key] for key in defaults} == defaults + assert cursor._executor._max_workers == 1 diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index d739bc22f..29d36948c 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -881,7 +881,7 @@ def test_duplicate_column_names_new_columns(self, polars_cursor): @pytest.mark.parametrize( "execute_kwargs", - [{}, {"block_size": 2048, "cache_type": "none", "max_workers": 3, "chunksize": 20}], + [{}, {"block_size": 2048, "cache_type": "none", "s3_max_workers": 3, "chunksize": 20}], ) def test_read_options(self, execute_kwargs): """The cursor's read options reach the result set, and execute() overrides them. @@ -891,7 +891,7 @@ def test_read_options(self, execute_kwargs): cursor_kwargs = { "block_size": 1024, "cache_type": "bytes", - "max_workers": 2, + "s3_max_workers": 2, "chunksize": 10, } cursor = PolarsCursor( @@ -908,6 +908,14 @@ def test_read_options(self, execute_kwargs): patch("pyathena.polars.cursor.AthenaPolarsResultSet") as result_set_class, ): cursor.execute("SELECT 1", **execute_kwargs) - kwargs = result_set_class.call_args.kwargs + first_kwargs = result_set_class.call_args.kwargs.copy() + cursor.execute("SELECT 1") + kwargs = first_kwargs expected = {**cursor_kwargs, **execute_kwargs} + expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected + + defaults = cursor_kwargs.copy() + defaults["max_workers"] = defaults.pop("s3_max_workers") + second_kwargs = result_set_class.call_args.kwargs + assert {key: second_kwargs[key] for key in defaults} == defaults diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index 1c5eed6fa..7bd061078 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -1,10 +1,21 @@ +from inspect import isawaitable from typing import Any +from unittest.mock import MagicMock, patch import pytest import tenacity.nap from botocore.exceptions import ClientError from pyathena import DataError +from pyathena.aio.arrow.cursor import AioArrowCursor +from pyathena.aio.pandas.cursor import AioPandasCursor +from pyathena.aio.polars.cursor import AioPolarsCursor +from pyathena.arrow.async_cursor import AsyncArrowCursor +from pyathena.arrow.cursor import ArrowCursor +from pyathena.pandas.async_cursor import AsyncPandasCursor +from pyathena.pandas.cursor import PandasCursor +from pyathena.polars.async_cursor import AsyncPolarsCursor +from pyathena.polars.cursor import PolarsCursor from pyathena.util import ( RetryConfig, _is_throttling_error, @@ -372,3 +383,74 @@ def call(): assert caught.value is error assert calls == expected_calls + + +S3_WORKER_CURSORS = ( + PandasCursor, + AsyncPandasCursor, + AioPandasCursor, + PolarsCursor, + AsyncPolarsCursor, + AioPolarsCursor, + ArrowCursor, + AsyncArrowCursor, + AioArrowCursor, +) + + +class TestS3WorkerOptions: + @staticmethod + async def _execute(cursor, kwargs): + result = cursor.execute("SELECT 1", **kwargs) + if isawaitable(result): + await result + + @staticmethod + def _cursor(cursor_class, **kwargs): + return cursor_class( + connection=MagicMock(), + converter=MagicMock(), + formatter=MagicMock(), + retry_config=RetryConfig(), + **kwargs, + ) + + @pytest.mark.parametrize("cursor_class", S3_WORKER_CURSORS) + @pytest.mark.parametrize( + ("value", "error"), [(0, ValueError), (-1, ValueError), (True, TypeError), (1.5, TypeError)] + ) + def test_s3_workers_invalid_constructor(self, cursor_class, value, error): + with pytest.raises(error, match="s3_max_workers"): + self._cursor(cursor_class, s3_max_workers=value) + + @pytest.mark.parametrize("cursor_class", S3_WORKER_CURSORS) + @pytest.mark.parametrize( + "kwargs", [{"s3_max_workers": 0}, {"s3_max_workers": -1}, {"max_workers": 2}] + ) + async def test_s3_workers_invalid_execute(self, cursor_class, kwargs): + cursor = self._cursor(cursor_class) + error = TypeError if "max_workers" in kwargs else ValueError + try: + with patch.object(cursor, "_execute") as execute: + with pytest.raises(error, match="s3_max_workers"): + await self._execute(cursor, kwargs) + execute.assert_not_called() + finally: + closed = cursor.close() + if isawaitable(closed): + await closed + + @pytest.mark.parametrize( + "cursor_class", + [PandasCursor, PolarsCursor, ArrowCursor, AioPandasCursor, AioPolarsCursor, AioArrowCursor], + ) + def test_s3_workers_reject_former_constructor_keyword(self, cursor_class): + with pytest.raises(TypeError, match="s3_max_workers"): + self._cursor(cursor_class, max_workers=2) + + @pytest.mark.parametrize( + "cursor_class", [PandasCursor, PolarsCursor, AioPandasCursor, AioPolarsCursor] + ) + def test_s3_workers_none_is_arrow_only(self, cursor_class): + with pytest.raises(TypeError, match="s3_max_workers"): + self._cursor(cursor_class, s3_max_workers=None) From d30dff2ce1e35358c26b75e70fb42b2a89262d14 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 6 Oct 2026 01:34:43 +0900 Subject: [PATCH 2/6] Migrate benchmark worker options and serialize Arrow S3 client creation --- benchmarks/pyathena_bench/adapters.py | 12 +++---- benchmarks/tests/test_adapters.py | 11 ++++-- pyathena/arrow/result_set.py | 13 +++---- tests/pyathena/arrow/test_result_set.py | 46 +++++++++++++++++++++++++ 4 files changed, 66 insertions(+), 16 deletions(-) diff --git a/benchmarks/pyathena_bench/adapters.py b/benchmarks/pyathena_bench/adapters.py index 96263b2c5..e452722d0 100644 --- a/benchmarks/pyathena_bench/adapters.py +++ b/benchmarks/pyathena_bench/adapters.py @@ -179,8 +179,10 @@ def connection(self) -> Any: def cursor(self, connection: Any) -> Any: case = self.case kwargs: dict[str, Any] = {"arraysize": case.arraysize} - if case.api == "thread" or case.family in {"pandas", "polars"}: + if case.api == "thread": kwargs["max_workers"] = self.settings.executor_workers + if case.family in {"pandas", "polars"}: + kwargs["s3_max_workers"] = self.settings.executor_workers if case.family in {"pandas", "polars", "arrow"}: kwargs["unload"] = case.transport == "unload" if case.family in {"pandas", "polars"}: @@ -251,14 +253,8 @@ async def measure_query(adapter: Adapter, cursor: Any, sql: str, expected: int) count = consume(result, case) query_id = None # The botocore observer captures IDs even for iterators. elif case.api == "thread": - # AsyncPandasCursor forwards result-reader options through execute(). - options = ( - {"max_workers": adapter.settings.executor_workers} - if case.family == "pandas" - else {} - ) query_id, future = await asyncio.to_thread( - cursor.execute, sql, cache_size=0, result_reuse_enable=False, **options + cursor.execute, sql, cache_size=0, result_reuse_enable=False ) result = await asyncio.wrap_future(future) ready = time.perf_counter() diff --git a/benchmarks/tests/test_adapters.py b/benchmarks/tests/test_adapters.py index 76b738f0b..f09ff8624 100644 --- a/benchmarks/tests/test_adapters.py +++ b/benchmarks/tests/test_adapters.py @@ -129,13 +129,20 @@ def unexpected_request(*args, **kwargs): session = boto3.Session( aws_access_key_id="testing", aws_secret_access_key="testing", region_name="us-west-2" ) - for case in matrix(Settings(), "single", "flat"): + settings = Settings(executor_workers=3) + for case in matrix(settings, "single", "flat"): if case.family == "wrangler" or case.unsupported: continue - adapter = Adapter(Settings(), case, session, "scratch", "s3://out/", "temp") + adapter = Adapter(settings, case, session, "scratch", "s3://out/", "temp") with adapter.connection() as connection: cursor = adapter.cursor(connection) assert cursor.arraysize == case.arraysize + if case.api == "thread": + assert cursor._executor._max_workers == 3 + if case.family in {"pandas", "polars"}: + assert cursor._s3_max_workers == 3 + elif case.family == "arrow": + assert cursor._s3_max_workers is None cursor.close() @pytest.mark.parametrize("api", ["sync", "thread", "aio"]) diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 8fe01030f..43b72c770 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -182,12 +182,13 @@ def _create_s3_file_system(self): overrides["connect_timeout"] = self._connect_timeout if self._request_timeout is not None: overrides["read_timeout"] = self._request_timeout - client = connection.session.client( - "s3", - region_name=connection.region_name, - config=connection.s3_config.merge(Config(**overrides)), - **connection._s3_client_kwargs, - ) + with connection._s3_client_lock: + client = connection.session.client( + "s3", + region_name=connection.region_name, + config=connection.s3_config.merge(Config(**overrides)), + **connection._s3_client_kwargs, + ) self._s3_resources.callback(client.close) filesystem = S3FileSystem( connection=connection, diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index cfadf3d2e..e342acf0d 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -4,6 +4,8 @@ # See LICENSE or https://opensource.org/licenses/MIT. # # SPDX-License-Identifier: MIT +import threading +from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime from io import BytesIO from unittest.mock import MagicMock, patch @@ -19,6 +21,7 @@ from pyathena.arrow.converter import DefaultArrowTypeConverter from pyathena.arrow.result_set import AthenaArrowResultSet +from pyathena.connection import Connection from pyathena.filesystem.s3_executor import S3ThreadPoolExecutor from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet @@ -201,3 +204,46 @@ def test_s3_workers_native_default(self): access_key="dummy", secret_key="dummy", session_token=None, region="us-east-1" ) connection.session.client.assert_not_called() + + def test_s3_workers_serializes_client_creation(self): + """Dedicated and shared S3 clients use the same creation lock.""" + session = MagicMock() + connection = Connection( + session=session, region_name="us-east-1", s3_staging_dir="s3://bucket/path/" + ) + created = [] + + def create_client(*args, **kwargs): + assert args == ("s3",) + assert connection._s3_client_lock.locked() + client = MagicMock() + created.append(client) + return client + + session.client.side_effect = create_client + barrier = threading.Barrier(4) + + def create_result(index): + barrier.wait(timeout=5) + if index == 0: + return connection.s3_client + return AthenaArrowResultSet( + connection=connection, + converter=DefaultArrowTypeConverter(), + query_execution=MagicMock(state=AthenaQueryExecution.STATE_FAILED), + arraysize=1, + retry_config=RetryConfig(), + s3_max_workers=2, + request_timeout=4.5, + ) + + with ThreadPoolExecutor(max_workers=4) as executor: + results = list(executor.map(create_result, range(4))) + assert len(created) == 4 + shared_client = results[0] + for client in created: + if client is shared_client: + client.close.assert_not_called() + else: + client.close.assert_called_once_with() + assert not connection._s3_client_lock.locked() From 8f1d1afc2db345bb56ac31b6c248001cd8afa973 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 9 Oct 2026 23:17:13 +0900 Subject: [PATCH 3/6] Use standalone functions for S3 worker option tests --- tests/pyathena/test_util.py | 108 ++++++++++++++++++------------------ 1 file changed, 55 insertions(+), 53 deletions(-) diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index 7bd061078..187dfa272 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -398,59 +398,61 @@ def call(): ) -class TestS3WorkerOptions: - @staticmethod - async def _execute(cursor, kwargs): - result = cursor.execute("SELECT 1", **kwargs) - if isawaitable(result): - await result - - @staticmethod - def _cursor(cursor_class, **kwargs): - return cursor_class( - connection=MagicMock(), - converter=MagicMock(), - formatter=MagicMock(), - retry_config=RetryConfig(), - **kwargs, - ) - - @pytest.mark.parametrize("cursor_class", S3_WORKER_CURSORS) - @pytest.mark.parametrize( - ("value", "error"), [(0, ValueError), (-1, ValueError), (True, TypeError), (1.5, TypeError)] +async def _execute_s3_worker_query(cursor, kwargs): + result = cursor.execute("SELECT 1", **kwargs) + if isawaitable(result): + await result + + +def _s3_worker_cursor(cursor_class, **kwargs): + return cursor_class( + connection=MagicMock(), + converter=MagicMock(), + formatter=MagicMock(), + retry_config=RetryConfig(), + **kwargs, ) - def test_s3_workers_invalid_constructor(self, cursor_class, value, error): - with pytest.raises(error, match="s3_max_workers"): - self._cursor(cursor_class, s3_max_workers=value) - @pytest.mark.parametrize("cursor_class", S3_WORKER_CURSORS) - @pytest.mark.parametrize( - "kwargs", [{"s3_max_workers": 0}, {"s3_max_workers": -1}, {"max_workers": 2}] - ) - async def test_s3_workers_invalid_execute(self, cursor_class, kwargs): - cursor = self._cursor(cursor_class) - error = TypeError if "max_workers" in kwargs else ValueError - try: - with patch.object(cursor, "_execute") as execute: - with pytest.raises(error, match="s3_max_workers"): - await self._execute(cursor, kwargs) - execute.assert_not_called() - finally: - closed = cursor.close() - if isawaitable(closed): - await closed - - @pytest.mark.parametrize( - "cursor_class", - [PandasCursor, PolarsCursor, ArrowCursor, AioPandasCursor, AioPolarsCursor, AioArrowCursor], - ) - def test_s3_workers_reject_former_constructor_keyword(self, cursor_class): - with pytest.raises(TypeError, match="s3_max_workers"): - self._cursor(cursor_class, max_workers=2) - @pytest.mark.parametrize( - "cursor_class", [PandasCursor, PolarsCursor, AioPandasCursor, AioPolarsCursor] - ) - def test_s3_workers_none_is_arrow_only(self, cursor_class): - with pytest.raises(TypeError, match="s3_max_workers"): - self._cursor(cursor_class, s3_max_workers=None) +@pytest.mark.parametrize("cursor_class", S3_WORKER_CURSORS) +@pytest.mark.parametrize( + ("value", "error"), [(0, ValueError), (-1, ValueError), (True, TypeError), (1.5, TypeError)] +) +def test_s3_workers_invalid_constructor(cursor_class, value, error): + with pytest.raises(error, match="s3_max_workers"): + _s3_worker_cursor(cursor_class, s3_max_workers=value) + + +@pytest.mark.parametrize("cursor_class", S3_WORKER_CURSORS) +@pytest.mark.parametrize( + "kwargs", [{"s3_max_workers": 0}, {"s3_max_workers": -1}, {"max_workers": 2}] +) +async def test_s3_workers_invalid_execute(cursor_class, kwargs): + cursor = _s3_worker_cursor(cursor_class) + error = TypeError if "max_workers" in kwargs else ValueError + try: + with patch.object(cursor, "_execute") as execute: + with pytest.raises(error, match="s3_max_workers"): + await _execute_s3_worker_query(cursor, kwargs) + execute.assert_not_called() + finally: + closed = cursor.close() + if isawaitable(closed): + await closed + + +@pytest.mark.parametrize( + "cursor_class", + [PandasCursor, PolarsCursor, ArrowCursor, AioPandasCursor, AioPolarsCursor, AioArrowCursor], +) +def test_s3_workers_reject_former_constructor_keyword(cursor_class): + with pytest.raises(TypeError, match="s3_max_workers"): + _s3_worker_cursor(cursor_class, max_workers=2) + + +@pytest.mark.parametrize( + "cursor_class", [PandasCursor, PolarsCursor, AioPandasCursor, AioPolarsCursor] +) +def test_s3_workers_none_is_arrow_only(cursor_class): + with pytest.raises(TypeError, match="s3_max_workers"): + _s3_worker_cursor(cursor_class, s3_max_workers=None) From c120e3011a9d20dbf8f12e27765df0794da39973 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 01:04:13 +0900 Subject: [PATCH 4/6] Limit S3 worker settings to the pandas and Polars cursors Arrow cursors keep the native PyArrow S3 filesystem without an s3_max_workers option, so the dedicated S3 client, the S3FileSystem client argument, and their tests are removed. Former keyword checks and value validation before query submission are removed; unknown cursor keywords are handled by #1093, and reader options by the readers. The read option tests keep the renamed argument and the independent thread pool size. The 4.0 migration notes move to usage.md. Co-Authored-By: Claude Opus 5.5 --- benchmarks/tests/test_adapters.py | 2 - docs/aio.md | 2 - docs/arrow.md | 2 - docs/cursor.md | 45 ----- docs/pandas.md | 2 - docs/polars.md | 2 - docs/usage.md | 20 +++ pyathena/aio/arrow/cursor.py | 23 +-- pyathena/aio/pandas/cursor.py | 4 +- pyathena/aio/polars/cursor.py | 4 +- pyathena/arrow/async_cursor.py | 26 +-- pyathena/arrow/cursor.py | 26 +-- pyathena/arrow/result_set.py | 97 +++------- pyathena/filesystem/s3.py | 11 +- pyathena/pandas/async_cursor.py | 4 +- pyathena/pandas/cursor.py | 4 +- pyathena/polars/async_cursor.py | 4 +- pyathena/polars/cursor.py | 4 +- pyathena/util.py | 24 --- tests/pyathena/aio/arrow/test_cursor.py | 31 +--- tests/pyathena/aio/pandas/test_cursor.py | 9 +- tests/pyathena/aio/polars/test_cursor.py | 9 +- tests/pyathena/arrow/test_async_cursor.py | 35 +--- tests/pyathena/arrow/test_cursor.py | 31 +--- tests/pyathena/arrow/test_result_set.py | 195 +-------------------- tests/pyathena/filesystem/test_s3.py | 14 -- tests/pyathena/pandas/test_async_cursor.py | 10 +- tests/pyathena/pandas/test_cursor.py | 9 +- tests/pyathena/polars/test_async_cursor.py | 10 +- tests/pyathena/polars/test_cursor.py | 9 +- tests/pyathena/test_util.py | 84 --------- 31 files changed, 95 insertions(+), 657 deletions(-) diff --git a/benchmarks/tests/test_adapters.py b/benchmarks/tests/test_adapters.py index f09ff8624..5eb67dd7a 100644 --- a/benchmarks/tests/test_adapters.py +++ b/benchmarks/tests/test_adapters.py @@ -141,8 +141,6 @@ def unexpected_request(*args, **kwargs): assert cursor._executor._max_workers == 3 if case.family in {"pandas", "polars"}: assert cursor._s3_max_workers == 3 - elif case.family == "arrow": - assert cursor._s3_max_workers is None cursor.close() @pytest.mark.parametrize("api", ["sync", "thread", "aio"]) diff --git a/docs/aio.md b/docs/aio.md index 84bf0cab1..76d917d02 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -2,8 +2,6 @@ # Native Asyncio Cursors -See {ref}`cursor-workers` for query and S3 worker settings and the 4.0.0 migration. - PyAthena provides native asyncio cursor implementations under `pyathena.aio`. These cursors use `asyncio.sleep` for polling and `asyncio.to_thread` for boto3 calls, keeping the event loop free. Concurrency comes from asyncio tasks rather than a diff --git a/docs/arrow.md b/docs/arrow.md index 6f01ecf8b..5306f0902 100644 --- a/docs/arrow.md +++ b/docs/arrow.md @@ -2,8 +2,6 @@ # Arrow -See {ref}`cursor-workers` for query and S3 worker settings and the 4.0.0 migration. - (arrow-cursor)= ## ArrowCursor diff --git a/docs/cursor.md b/docs/cursor.md index 6f3c41fb7..a3db13a2a 100644 --- a/docs/cursor.md +++ b/docs/cursor.md @@ -1,50 +1,5 @@ # Cursor -(cursor-workers)= - -## Worker settings in 4.0.0 - -The thread-pool cursors use `max_workers` for their pool of tasks that wait for queries and collect results. -Query submission happens before the task is queued, so this setting does not cap the number of queries running in Athena. -The pandas, Arrow, and Polars cursors use `s3_max_workers` for the workers of each PyAthena S3 file reader. -It does not set a shared limit across queries or change the parsing library's CPU thread pool. - -| Cursor | `max_workers` | `s3_max_workers` default | -|--------|---------------|--------------------------| -| `PandasCursor`, `PolarsCursor`, `AioPandasCursor`, `AioPolarsCursor` | No query pool | `(cpu_count() or 1) * 5` | -| `AsyncPandasCursor`, `AsyncPolarsCursor` | Query task pool | `(cpu_count() or 1) * 5` | -| `ArrowCursor`, `AioArrowCursor` | No query pool | `None` (native PyArrow S3 filesystem) | -| `AsyncArrowCursor` | Query task pool | `None` (native PyArrow S3 filesystem) | - -Pass `s3_max_workers` to `execute()` to override the setting for one query. -The next call without an override uses the cursor's setting again. -For Arrow, a positive integer selects PyAthena's S3 filesystem through PyArrow's fsspec adapter for both CSV and UNLOAD results. -An explicit `None` selects the native PyArrow filesystem for that query. -On the PyAthena path, Arrow's `request_timeout` sets the boto3 read timeout; `connect_timeout` sets the connection timeout. -The connection's S3 configuration supplies defaults for unspecified timeouts. - -```python -from pyathena.pandas.async_cursor import AsyncPandasCursor - -cursor = connection.cursor(AsyncPandasCursor, max_workers=4, s3_max_workers=2) -query_id, future = cursor.execute("SELECT 1", s3_max_workers=3) -result = future.result() -``` - -The S3 setting applies when PyAthena's S3 filesystem reads the results. -User-provided pandas `filesystem` or `storage_options` can replace that filesystem and its settings. -Polars uses PyAthena's filesystem for CSV results; its native Parquet reader uses its own concurrency settings. - -### Migrate the former S3 argument - -This is a breaking API change in 4.0.0. -Replace `max_workers=N` with `s3_max_workers=N` in synchronous and native asyncio pandas/Polars cursor constructors, and in pandas/Polars `execute()` calls. -Keep `max_workers` in thread-pool cursor constructors to size the query task pool. -For `AsyncPolarsCursor`, specify both arguments to preserve the former behavior of controlling both pools with one value. -`AsyncPandasCursor` now also accepts an S3 worker default in its constructor. -The former S3 keyword raises `TypeError`; nonpositive S3 worker values raise `ValueError` before a query starts. -The `max_workers` names on result sets, S3 filesystems, and pandas upload utilities retain their existing meanings. - (default_cursor)= ## DefaultCursor diff --git a/docs/pandas.md b/docs/pandas.md index 5e43c33b5..f4816ff49 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -2,8 +2,6 @@ # Pandas -See {ref}`cursor-workers` for query and S3 worker settings and the 4.0.0 migration. - (as-dataframe)= ## As DataFrame diff --git a/docs/polars.md b/docs/polars.md index 8fedba7d0..f347e4214 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -2,8 +2,6 @@ # Polars -See {ref}`cursor-workers` for query and S3 worker settings and the 4.0.0 migration. - (polars-cursor)= ## PolarsCursor diff --git a/docs/usage.md b/docs/usage.md index 4a2febcb6..d54bb3a12 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -40,6 +40,26 @@ cursor.execute("SELECT %(value)s", {"value": 1}, work_group="YOUR_WORK_GROUP") Existing keyword arguments retain their defaults and behavior. The cursor class can still be selected positionally with `connection.cursor(DictCursor)`. +## Worker settings in PyAthena 4.0 + +The pandas and Polars cursors take the number of threads that read each result file from S3 as `s3_max_workers`. +It can be set in the cursor constructor and overridden for one query in `execute()`. +The default is `(cpu_count() or 1) * 5`. +Polars reads UNLOAD results with its native Parquet reader, which does not use this setting. + +The thread-pool cursors, such as `AsyncCursor`, `AsyncPandasCursor`, and `AsyncPolarsCursor`, keep `max_workers` for the size of the thread pool that waits for queries and collects their results. + +```python +from pyathena.pandas.async_cursor import AsyncPandasCursor + +cursor = connection.cursor(AsyncPandasCursor, max_workers=4, s3_max_workers=2) +query_id, future = cursor.execute("SELECT 1", s3_max_workers=3) +``` + +Before PyAthena 4.0, `max_workers` set the S3 read workers of `PandasCursor`, `PolarsCursor`, `AioPandasCursor`, and `AioPolarsCursor`, and of `execute()` in the pandas and Polars cursors. +Replace it with `s3_max_workers` in those places. +`AsyncPolarsCursor` used its `max_workers` for both the thread pool and the S3 read workers; pass both arguments to keep the former values. + ## Managed query result storage When using a workgroup with [managed query result storage](https://docs.aws.amazon.com/athena/latest/ug/managed-results.html) enabled, diff --git a/pyathena/aio/arrow/cursor.py b/pyathena/aio/arrow/cursor.py index f91d220e0..c297dd0e1 100644 --- a/pyathena/aio/arrow/cursor.py +++ b/pyathena/aio/arrow/cursor.py @@ -17,7 +17,7 @@ from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions -from pyathena.util import _validate_s3_max_workers, override +from pyathena.util import override if TYPE_CHECKING: import polars as pl @@ -55,7 +55,6 @@ def __init__( result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, connect_timeout: float | None = None, request_timeout: float | None = None, - s3_max_workers: int | None = None, **kwargs, ) -> None: """Initialize an AioArrowCursor. @@ -73,17 +72,13 @@ def __init__( unload: Whether to wrap queries in ``UNLOAD`` and read the Parquet output. result_reuse_enable: Whether to enable Athena query result reuse. result_reuse_minutes: Maximum age of a reused query result in minutes. - connect_timeout: S3 connection timeout in seconds. If None, use the - selected filesystem's default: PyArrow's default on the native path, - or the connection's S3 configuration on the PyAthena path. - request_timeout: S3 request timeout on the native PyArrow path, or boto3 - read timeout on the PyAthena path. If None, use that path's default. - s3_max_workers: S3 read workers per file. If None, use the native - pyarrow filesystem; otherwise use PyAthena's S3 filesystem. + connect_timeout: Connection timeout in seconds of the pyarrow S3 filesystem + that reads the results. If None, the pyarrow default is used. + request_timeout: Request timeout in seconds of the pyarrow S3 filesystem + that reads the results. If None, the pyarrow default is used. **kwargs: Other cursor arguments, such as ``connection`` and ``arraysize``, passed to the parent ``__init__``. """ - _validate_s3_max_workers(s3_max_workers, kwargs, allow_none=True) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -97,7 +92,6 @@ def __init__( result_reuse_minutes=result_reuse_minutes, **kwargs, ) - self._s3_max_workers = s3_max_workers self._unload = unload self._connect_timeout = connect_timeout self._request_timeout = request_timeout @@ -150,8 +144,7 @@ async def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. ``s3_max_workers`` - overrides the cursor's S3 read setting for this query. + **kwargs: Additional execution parameters. ``block_size`` sets the read block size for this query, and ``connect_timeout`` and ``request_timeout`` override the cursor's values. @@ -159,9 +152,6 @@ async def execute( Self reference for method chaining. """ self._reset_state() - _validate_s3_max_workers( - kwargs.get("s3_max_workers", self._s3_max_workers), kwargs, allow_none=True - ) options = ExecuteOptions.resolve( options, work_group=work_group, @@ -197,7 +187,6 @@ async def execute( unload_location=unload_location, connect_timeout=kwargs.pop("connect_timeout", self._connect_timeout), request_timeout=kwargs.pop("request_timeout", self._request_timeout), - s3_max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), result_set_type_hints=options.result_set_type_hints, **kwargs, ) diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index c052fc882..cc7f259d3 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -22,7 +22,7 @@ DefaultPandasUnloadTypeConverter, ) from pyathena.pandas.result_set import AthenaPandasResultSet, PandasDataFrameIterator -from pyathena.util import _validate_s3_max_workers, override +from pyathena.util import override if TYPE_CHECKING: from pandas import DataFrame @@ -92,7 +92,6 @@ def __init__( **kwargs: Other cursor arguments, such as ``connection`` and ``arraysize``, passed to the parent ``__init__``. """ - _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -179,7 +178,6 @@ async def execute( Self reference for method chaining. """ self._reset_state() - _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 7f32cdf36..17a551a99 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -18,7 +18,7 @@ DefaultPolarsUnloadTypeConverter, ) from pyathena.polars.result_set import AthenaPolarsResultSet -from pyathena.util import _validate_s3_max_workers, override +from pyathena.util import override if TYPE_CHECKING: import polars as pl @@ -84,7 +84,6 @@ def __init__( **kwargs: Other cursor arguments, such as ``connection`` and ``arraysize``, passed to the parent ``__init__``. """ - _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -163,7 +162,6 @@ async def execute( Self reference for method chaining. """ self._reset_state() - _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/arrow/async_cursor.py b/pyathena/arrow/async_cursor.py index 4adb1884e..fe6c65814 100644 --- a/pyathena/arrow/async_cursor.py +++ b/pyathena/arrow/async_cursor.py @@ -17,7 +17,7 @@ from pyathena.common import CursorIterator from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions -from pyathena.util import _validate_s3_max_workers, override +from pyathena.util import override _logger = logging.getLogger(__name__) @@ -80,7 +80,6 @@ def __init__( result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, connect_timeout: float | None = None, request_timeout: float | None = None, - s3_max_workers: int | None = None, **kwargs, ) -> None: """Initialize an AsyncArrowCursor. @@ -99,14 +98,12 @@ def __init__( unload: Enable UNLOAD for high-performance Parquet output. result_reuse_enable: Enable Athena query result reuse. result_reuse_minutes: Minutes to reuse cached results. - connect_timeout: S3 connection timeout in seconds. If None, use the - selected filesystem's default: PyArrow's default on the native path, - or the connection's S3 configuration on the PyAthena path. - request_timeout: S3 request timeout in seconds on the native PyArrow path, - or boto3 read timeout on the PyAthena path. If None, use the selected - filesystem's default. - s3_max_workers: S3 read workers per file. If None, use the native - pyarrow filesystem; otherwise use PyAthena's S3 filesystem. + connect_timeout: Socket connection timeout in seconds for S3 operations. + Defaults to AWS SDK default (typically 1 second) if not specified. + request_timeout: Request timeout in seconds for S3 operations. + Defaults to AWS SDK default (typically 3 seconds) if not specified. + Increase this value if you experience timeout errors when using + role assumption with STS or have high latency to S3. **kwargs: Additional connection parameters. Example: @@ -117,7 +114,6 @@ def __init__( ... request_timeout=30.0 ... ) """ - _validate_s3_max_workers(s3_max_workers, kwargs, allow_none=True) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -133,7 +129,6 @@ def __init__( result_reuse_minutes=result_reuse_minutes, **kwargs, ) - self._s3_max_workers = s3_max_workers self._unload = unload self._connect_timeout = connect_timeout self._request_timeout = request_timeout @@ -179,7 +174,6 @@ def _collect_result_set( unload_location=unload_location, connect_timeout=kwargs.pop("connect_timeout", self._connect_timeout), request_timeout=kwargs.pop("request_timeout", self._request_timeout), - s3_max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), result_set_type_hints=result_set_type_hints, **kwargs, ) @@ -218,17 +212,13 @@ def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. ``s3_max_workers`` - overrides the cursor's S3 read setting for this query. + **kwargs: Additional execution parameters. ``block_size`` sets the read block size for this query, and ``connect_timeout`` and ``request_timeout`` override the cursor's values. Returns: Tuple of (query_id, future) where future resolves to AthenaArrowResultSet. """ - _validate_s3_max_workers( - kwargs.get("s3_max_workers", self._s3_max_workers), kwargs, allow_none=True - ) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/arrow/cursor.py b/pyathena/arrow/cursor.py index 4fe9798dc..e39c58237 100644 --- a/pyathena/arrow/cursor.py +++ b/pyathena/arrow/cursor.py @@ -16,7 +16,7 @@ from pyathena.model import AthenaQueryExecution from pyathena.options import ExecuteOptions from pyathena.result_set import WithFetch -from pyathena.util import _validate_s3_max_workers, override +from pyathena.util import override if TYPE_CHECKING: import polars as pl @@ -70,7 +70,6 @@ def __init__( result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, connect_timeout: float | None = None, request_timeout: float | None = None, - s3_max_workers: int | None = None, **kwargs, ) -> None: """Initialize an ArrowCursor. @@ -87,14 +86,12 @@ def __init__( unload: Enable UNLOAD for high-performance Parquet output. result_reuse_enable: Enable Athena query result reuse. result_reuse_minutes: Minutes to reuse cached results. - connect_timeout: S3 connection timeout in seconds. If None, use the - selected filesystem's default: PyArrow's default on the native path, - or the connection's S3 configuration on the PyAthena path. - request_timeout: S3 request timeout in seconds on the native PyArrow path, - or boto3 read timeout on the PyAthena path. If None, use the selected - filesystem's default. - s3_max_workers: S3 read workers per file. If None, use the native - pyarrow filesystem; otherwise use PyAthena's S3 filesystem. + connect_timeout: Socket connection timeout in seconds for S3 operations. + Defaults to AWS SDK default (typically 1 second) if not specified. + request_timeout: Request timeout in seconds for S3 operations. + Defaults to AWS SDK default (typically 3 seconds) if not specified. + Increase this value if you experience timeout errors when using + role assumption with STS or have high latency to S3. **kwargs: Additional connection parameters. Example: @@ -105,7 +102,6 @@ def __init__( ... request_timeout=30 ... ) """ - _validate_s3_max_workers(s3_max_workers, kwargs, allow_none=True) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -119,7 +115,6 @@ def __init__( result_reuse_minutes=result_reuse_minutes, **kwargs, ) - self._s3_max_workers = s3_max_workers self._unload = unload self._connect_timeout = connect_timeout self._request_timeout = request_timeout @@ -175,8 +170,7 @@ def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. ``s3_max_workers`` - overrides the cursor's S3 read setting for this query. + **kwargs: Additional execution parameters. ``block_size`` sets the read block size for this query, and ``connect_timeout`` and ``request_timeout`` override the cursor's values. @@ -188,9 +182,6 @@ def execute( >>> table = cursor.as_arrow() # Returns Apache Arrow Table """ self._reset_state() - _validate_s3_max_workers( - kwargs.get("s3_max_workers", self._s3_max_workers), kwargs, allow_none=True - ) options = ExecuteOptions.resolve( options, work_group=work_group, @@ -224,7 +215,6 @@ def execute( unload_location=unload_location, connect_timeout=kwargs.pop("connect_timeout", self._connect_timeout), request_timeout=kwargs.pop("request_timeout", self._request_timeout), - s3_max_workers=kwargs.pop("s3_max_workers", self._s3_max_workers), result_set_type_hints=options.result_set_type_hints, **kwargs, ) diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 43b72c770..2b1b9cbb9 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -4,23 +4,19 @@ import logging from collections.abc import Callable -from contextlib import ExitStack from typing import ( TYPE_CHECKING, Any, ClassVar, ) -from botocore.config import Config - from pyathena import OperationalError from pyathena.arrow.converter import _to_timestamp from pyathena.arrow.util import to_column_info from pyathena.converter import Converter, _to_default -from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet -from pyathena.util import RetryConfig, _validate_s3_max_workers, override, parse_output_location +from pyathena.util import RetryConfig, override, parse_output_location if TYPE_CHECKING: import polars as pl @@ -98,8 +94,6 @@ def __init__( connect_timeout: float | None = None, request_timeout: float | None = None, result_set_type_hints: dict[str | int, str] | None = None, - *, - s3_max_workers: int | None = None, **kwargs, ) -> None: """Initialize the result set and load the query results into an Arrow Table. @@ -117,21 +111,16 @@ def __init__( instead of the CSV results. unload_location: The S3 location of the ``UNLOAD`` output. If None, it is derived from the first file in the data manifest. - connect_timeout: Connection timeout in seconds for the selected S3 filesystem. - request_timeout: Native PyArrow request timeout, or boto3 read timeout - on the PyAthena path, in seconds. + connect_timeout: The connect timeout in seconds for the pyarrow S3 filesystem. + request_timeout: The request timeout in seconds for the pyarrow S3 filesystem. result_set_type_hints: Athena type signatures for complex-type columns, keyed by column name (case-insensitive) or zero-based column index. - s3_max_workers: S3 read workers per file. None uses the native pyarrow - filesystem. A positive integer uses PyAthena's S3 filesystem; - request_timeout then sets the boto3 read timeout. **kwargs: Additional keyword arguments, stored but not used. Raises: ProgrammingError: If ``query_execution`` is not given. OperationalError: If reading the query results fails. """ - _validate_s3_max_workers(s3_max_workers, kwargs, allow_none=True) super().__init__( connection=connection, converter=converter, @@ -148,24 +137,21 @@ def __init__( self._connect_timeout = connect_timeout self._request_timeout = request_timeout self._kwargs = kwargs - self._s3_max_workers = s3_max_workers - self._s3_resources = ExitStack() - with self._s3_resources: - self._fs = self._create_s3_file_system() - if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location: - self._table = self._as_arrow() - elif self.state == AthenaQueryExecution.STATE_SUCCEEDED: - # Without a result file, as with managed query result storage, the rows from - # GetQueryResults are read as a CSV result file. - self._table = self._read_csv() - else: - import pyarrow as pa + self._fs = self._create_s3_file_system() + if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location: + self._table = self._as_arrow() + elif self.state == AthenaQueryExecution.STATE_SUCCEEDED: + # Without a result file, as with managed query result storage, the rows from + # GetQueryResults are read as a CSV result file. + self._table = self._read_csv() + else: + import pyarrow as pa - self._table = pa.Table.from_pydict({}) + self._table = pa.Table.from_pydict({}) self._batches = iter(self._table.to_batches(arraysize)) def _create_s3_file_system(self): - """Create the selected S3 filesystem from the connection settings. + """Create a pyarrow ``S3FileSystem`` from the connection settings. Returns: The pyarrow S3 filesystem for reading the query results. @@ -174,30 +160,6 @@ def _create_s3_file_system(self): connection = self.connection - if self._s3_max_workers is not None: - client = None - if self._connect_timeout is not None or self._request_timeout is not None: - overrides = {} - if self._connect_timeout is not None: - overrides["connect_timeout"] = self._connect_timeout - if self._request_timeout is not None: - overrides["read_timeout"] = self._request_timeout - with connection._s3_client_lock: - client = connection.session.client( - "s3", - region_name=connection.region_name, - config=connection.s3_config.merge(Config(**overrides)), - **connection._s3_client_kwargs, - ) - self._s3_resources.callback(client.close) - filesystem = S3FileSystem( - connection=connection, - max_workers=self._s3_max_workers, - s3_client=client, - skip_instance_cache=True, - ) - return fs.PyFileSystem(fs.FSSpecHandler(filesystem)) - # Build timeout parameters dict timeout_kwargs = {} if self._connect_timeout is not None: @@ -406,23 +368,20 @@ def _read_csv(self) -> Table: return pa.Table.from_pydict({}) try: - with ( - self._fs.open_input_stream(location) if data is None else pa.BufferReader(data) - ) as stream: - table = csv.read_csv( - stream, - read_options=read_opts, - parse_options=parse_opts, - convert_options=csv.ConvertOptions( - strings_can_be_null=bool(binary_columns), - quoted_strings_can_be_null=False, - timestamp_parsers=self.timestamp_parsers, - column_types={ - **column_types, - **{column_names[i]: pa.string() for i in timestamp_types}, - }, - ), - ) + table = csv.read_csv( + self._fs.open_input_stream(location) if data is None else pa.BufferReader(data), + read_options=read_opts, + parse_options=parse_opts, + convert_options=csv.ConvertOptions( + strings_can_be_null=bool(binary_columns), + quoted_strings_can_be_null=False, + timestamp_parsers=self.timestamp_parsers, + column_types={ + **column_types, + **{column_names[i]: pa.string() for i in timestamp_types}, + }, + ), + ) for index, type_ in timestamp_types.items(): table = table.set_column( index, diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 4ab7b2061..d3cf4537c 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -174,7 +174,6 @@ def __init__( allow_bucket_deletion: bool = False, version_aware: bool = False, *args, - s3_client: BaseClient | None = None, **kwargs, ) -> None: """Create a filesystem for Amazon S3. @@ -198,9 +197,6 @@ def __init__( allow_bucket_deletion: Whether ``rmdir`` may delete a bucket. version_aware: Whether reads pin the object version observed at open time. - s3_client: Optional S3 client overriding the connection or client - factory. The caller owns this client; the filesystem uses - the connection's retry policy when a connection is provided. *args: Passed to ``fsspec.AbstractFileSystem``. **kwargs: Passed to ``fsspec.AbstractFileSystem``; without a ``connection``, also s3fs-compatible client arguments. @@ -208,13 +204,12 @@ def __init__( the operations that accept ``RequestPayer``. """ super().__init__(*args, **kwargs) - if s3_client is not None: - client = s3_client - elif connection: + if connection: client = connection.s3_client + retry_config = connection.retry_config else: client = self._get_client_compatible_with_s3fs(**kwargs) - retry_config = connection.retry_config if connection else RetryConfig() + retry_config = RetryConfig() self.default_block_size = ( default_block_size if default_block_size else self.DEFAULT_BLOCK_SIZE ) diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index d49906529..4337065d5 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -18,7 +18,7 @@ DefaultPandasUnloadTypeConverter, ) from pyathena.pandas.result_set import AthenaPandasResultSet -from pyathena.util import _validate_s3_max_workers, override +from pyathena.util import override _logger = logging.getLogger(__name__) @@ -115,7 +115,6 @@ def __init__( **kwargs: Other cursor arguments, such as ``connection`` and ``converter``, passed to ``AsyncCursor.__init__``. """ - _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -247,7 +246,6 @@ def execute( Returns: Tuple of (query_id, future) where future resolves to AthenaPandasResultSet. """ - _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 0d542bf6b..37b68c044 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -21,7 +21,7 @@ ) from pyathena.pandas.result_set import AthenaPandasResultSet, PandasDataFrameIterator from pyathena.result_set import WithFetch -from pyathena.util import _validate_s3_max_workers, override +from pyathena.util import override if TYPE_CHECKING: from pandas import DataFrame @@ -115,7 +115,6 @@ def __init__( ``converter``, ``formatter``, and ``retry_config``. Pass pandas ``read_csv``/``read_parquet`` options to ``execute()`` instead. """ - _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -210,7 +209,6 @@ def execute( >>> df = cursor.as_pandas() # Returns pandas DataFrame """ self._reset_state() - _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index 85bf660ea..8f015621e 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -17,7 +17,7 @@ DefaultPolarsUnloadTypeConverter, ) from pyathena.polars.result_set import AthenaPolarsResultSet -from pyathena.util import _validate_s3_max_workers, override +from pyathena.util import override _logger = logging.getLogger(__name__) @@ -110,7 +110,6 @@ def __init__( >>> # With chunked processing >>> cursor = connection.cursor(AsyncPolarsCursor, chunksize=50000) """ - _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -249,7 +248,6 @@ def execute( >>> result_set = future.result() >>> df = result_set.as_polars() # Returns Polars DataFrame """ - _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 5e5811636..692c31d2a 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -21,7 +21,7 @@ ) from pyathena.polars.result_set import AthenaPolarsResultSet from pyathena.result_set import WithFetch -from pyathena.util import _validate_s3_max_workers, override +from pyathena.util import override if TYPE_CHECKING: import polars as pl @@ -112,7 +112,6 @@ def __init__( >>> # With chunked processing >>> cursor = connection.cursor(PolarsCursor, chunksize=50000) """ - _validate_s3_max_workers(s3_max_workers, kwargs) super().__init__( s3_staging_dir=s3_staging_dir, schema_name=schema_name, @@ -205,7 +204,6 @@ def execute( >>> df = cursor.as_polars() # Returns Polars DataFrame """ self._reset_state() - _validate_s3_max_workers(kwargs.get("s3_max_workers", self._s3_max_workers), kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/util.py b/pyathena/util.py index 05e5b4f61..a88dc6bf5 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -54,30 +54,6 @@ def override(func): ) -def _validate_s3_max_workers( - value: int | None, kwargs: dict[str, Any], *, allow_none: bool = False -) -> None: - """Validate cursor S3 workers and reject the former S3 keyword. - - Args: - value: The S3 worker setting to validate. - kwargs: Cursor constructor or execution keyword arguments. - allow_none: Whether None selects the native Arrow filesystem. - - Raises: - TypeError: If max_workers is used as a read option, or value is not an integer. - ValueError: If value is not positive. - """ - if "max_workers" in kwargs: - raise TypeError("Use s3_max_workers for S3 reads; max_workers configures async cursors.") - if value is None and allow_none: - return - if not isinstance(value, int) or isinstance(value, bool): - raise TypeError("s3_max_workers must be a positive integer") - if value <= 0: - raise ValueError("s3_max_workers must be greater than 0") - - def parse_output_location(output_location: str) -> tuple[str, str]: """Parse an S3 output location URL into bucket and key components. diff --git a/tests/pyathena/aio/arrow/test_cursor.py b/tests/pyathena/aio/arrow/test_cursor.py index 67c611835..fa4e14e11 100644 --- a/tests/pyathena/aio/arrow/test_cursor.py +++ b/tests/pyathena/aio/arrow/test_cursor.py @@ -19,20 +19,6 @@ class TestAioArrowCursor: - @pytest.mark.parametrize( - "aio_arrow_cursor", - [ - {"cursor_kwargs": {"s3_max_workers": 2, "unload": False}}, - {"cursor_kwargs": {"s3_max_workers": 2, "unload": True}}, - ], - indirect=True, - ) - async def test_s3_workers_read_results(self, aio_arrow_cursor): - """CSV and Parquet results use the configured PyAthena S3 reader.""" - await aio_arrow_cursor.execute("SELECT 1 AS value", s3_max_workers=1) - assert await aio_arrow_cursor.fetchall() == [(1,)] - assert aio_arrow_cursor.result_set._fs.handler.fs.max_workers == 1 - async def test_binary_null_vs_empty(self, aio_arrow_cursor): query = """SELECT * FROM (VALUES (1, CAST(NULL AS VARBINARY), 'null', CAST(NULL AS VARCHAR)), @@ -164,19 +150,14 @@ async def test_as_arrow_unload(self, aio_arrow_cursor): assert table.num_rows == 1 @pytest.mark.parametrize( - "execute_kwargs", - [ - {}, - {"connect_timeout": 3.0, "request_timeout": 4.0, "s3_max_workers": 3}, - {"s3_max_workers": None}, - ], + "execute_kwargs", [{}, {"connect_timeout": 3.0, "request_timeout": 4.0}] ) async def test_read_options(self, execute_kwargs): """The cursor's read options reach the result set, and execute() overrides them. No AWS calls; the query and its result set are mocked. """ - cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0, "s3_max_workers": 2} + cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0} query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) cursor = AioArrowCursor( connection=MagicMock(), @@ -191,12 +172,6 @@ async def test_read_options(self, execute_kwargs): patch("pyathena.aio.arrow.cursor.AthenaArrowResultSet") as result_set_class, ): await cursor.execute("SELECT 1", **execute_kwargs) - first_kwargs = result_set_class.call_args.kwargs.copy() - await cursor.execute("SELECT 1") - kwargs = first_kwargs + kwargs = result_set_class.call_args.kwargs expected = {**cursor_kwargs, **execute_kwargs} assert {key: kwargs[key] for key in expected} == expected - - defaults = cursor_kwargs.copy() - second_kwargs = result_set_class.call_args.kwargs - assert {key: second_kwargs[key] for key in defaults} == defaults diff --git a/tests/pyathena/aio/pandas/test_cursor.py b/tests/pyathena/aio/pandas/test_cursor.py index eebfb8778..5fd863a62 100644 --- a/tests/pyathena/aio/pandas/test_cursor.py +++ b/tests/pyathena/aio/pandas/test_cursor.py @@ -177,14 +177,7 @@ async def test_read_options(self, execute_kwargs): patch("pyathena.aio.pandas.cursor.AthenaPandasResultSet") as result_set_class, ): await cursor.execute("SELECT 1", **execute_kwargs) - first_kwargs = result_set_class.call_args.kwargs.copy() - await cursor.execute("SELECT 1") - kwargs = first_kwargs + kwargs = result_set_class.call_args.kwargs expected = {**cursor_kwargs, **execute_kwargs} expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected - - defaults = cursor_kwargs.copy() - defaults["max_workers"] = defaults.pop("s3_max_workers") - second_kwargs = result_set_class.call_args.kwargs - assert {key: second_kwargs[key] for key in defaults} == defaults diff --git a/tests/pyathena/aio/polars/test_cursor.py b/tests/pyathena/aio/polars/test_cursor.py index 049bd38bf..9c9003c74 100644 --- a/tests/pyathena/aio/polars/test_cursor.py +++ b/tests/pyathena/aio/polars/test_cursor.py @@ -163,14 +163,7 @@ async def test_read_options(self, execute_kwargs): patch("pyathena.aio.polars.cursor.AthenaPolarsResultSet") as result_set_class, ): await cursor.execute("SELECT 1", **execute_kwargs) - first_kwargs = result_set_class.call_args.kwargs.copy() - await cursor.execute("SELECT 1") - kwargs = first_kwargs + kwargs = result_set_class.call_args.kwargs expected = {**cursor_kwargs, **execute_kwargs} expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected - - defaults = cursor_kwargs.copy() - defaults["max_workers"] = defaults.pop("s3_max_workers") - second_kwargs = result_set_class.call_args.kwargs - assert {key: second_kwargs[key] for key in defaults} == defaults diff --git a/tests/pyathena/arrow/test_async_cursor.py b/tests/pyathena/arrow/test_async_cursor.py index 5d4d9cb09..30ae8f5ce 100644 --- a/tests/pyathena/arrow/test_async_cursor.py +++ b/tests/pyathena/arrow/test_async_cursor.py @@ -18,21 +18,6 @@ class TestAsyncArrowCursor: - @pytest.mark.parametrize( - "async_arrow_cursor", - [ - {"cursor_kwargs": {"s3_max_workers": 2, "unload": False}}, - {"cursor_kwargs": {"s3_max_workers": 2, "unload": True}}, - ], - indirect=True, - ) - def test_s3_workers_read_results(self, async_arrow_cursor): - """CSV and Parquet results use the configured PyAthena S3 reader.""" - _, future = async_arrow_cursor.execute("SELECT 1 AS value", s3_max_workers=1) - result = future.result() - assert result.fetchall() == [(1,)] - assert result._fs.handler.fs.max_workers == 1 - def test_binary_null_vs_empty(self, async_arrow_cursor): query = """SELECT * FROM (VALUES (1, CAST(NULL AS VARBINARY), 'null', CAST(NULL AS VARCHAR)), @@ -342,19 +327,14 @@ def test_empty_result_unload(self, async_arrow_cursor): assert table.shape[1] == 0 @pytest.mark.parametrize( - "execute_kwargs", - [ - {}, - {"connect_timeout": 3.0, "request_timeout": 4.0, "s3_max_workers": 3}, - {"s3_max_workers": None}, - ], + "execute_kwargs", [{}, {"connect_timeout": 3.0, "request_timeout": 4.0}] ) def test_read_options(self, execute_kwargs): """The cursor's read options reach the result set, and execute() overrides them. No AWS calls; the query and its result set are mocked. """ - cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0, "s3_max_workers": 2} + cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0} query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) with ( AsyncArrowCursor( @@ -362,7 +342,6 @@ def test_read_options(self, execute_kwargs): converter=MagicMock(), formatter=MagicMock(), retry_config=RetryConfig(), - max_workers=1, **cursor_kwargs, ) as cursor, patch.object(AsyncArrowCursor, "_execute", return_value="query_id"), @@ -371,14 +350,6 @@ def test_read_options(self, execute_kwargs): ): _, future = cursor.execute("SELECT 1", **execute_kwargs) future.result() - first_kwargs = result_set_class.call_args.kwargs.copy() - _, future = cursor.execute("SELECT 1") - future.result() - kwargs = first_kwargs + kwargs = result_set_class.call_args.kwargs expected = {**cursor_kwargs, **execute_kwargs} assert {key: kwargs[key] for key in expected} == expected - - defaults = cursor_kwargs.copy() - second_kwargs = result_set_class.call_args.kwargs - assert {key: second_kwargs[key] for key in defaults} == defaults - assert cursor._executor._max_workers == 1 diff --git a/tests/pyathena/arrow/test_cursor.py b/tests/pyathena/arrow/test_cursor.py index 49196e459..54352fa9c 100644 --- a/tests/pyathena/arrow/test_cursor.py +++ b/tests/pyathena/arrow/test_cursor.py @@ -35,20 +35,6 @@ class TestArrowCursor: - @pytest.mark.parametrize( - "arrow_cursor", - [ - {"cursor_kwargs": {"s3_max_workers": 2, "unload": False}}, - {"cursor_kwargs": {"s3_max_workers": 2, "unload": True}}, - ], - indirect=True, - ) - def test_s3_workers_read_results(self, arrow_cursor): - """CSV and Parquet results use the configured PyAthena S3 reader.""" - arrow_cursor.execute("SELECT 1 AS value", s3_max_workers=1) - assert arrow_cursor.fetchall() == [(1,)] - assert arrow_cursor.result_set._fs.handler.fs.max_workers == 1 - def test_binary_null_vs_empty(self, arrow_cursor): # The text column has the name of the binary column, and keeps its own NULL. query = """SELECT id, value, label, text_value AS value FROM (VALUES @@ -1140,19 +1126,14 @@ def test_duplicate_column_names_with_newline(self, arrow_cursor): assert arrow_cursor.as_arrow().column_names == ["a\nb", "a\nb"] @pytest.mark.parametrize( - "execute_kwargs", - [ - {}, - {"connect_timeout": 3.0, "request_timeout": 4.0, "s3_max_workers": 3}, - {"s3_max_workers": None}, - ], + "execute_kwargs", [{}, {"connect_timeout": 3.0, "request_timeout": 4.0}] ) def test_read_options(self, execute_kwargs): """The cursor's read options reach the result set, and execute() overrides them. No AWS calls; the query and its result set are mocked. """ - cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0, "s3_max_workers": 2} + cursor_kwargs = {"connect_timeout": 1.0, "request_timeout": 2.0} query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) cursor = ArrowCursor( connection=MagicMock(), @@ -1167,12 +1148,6 @@ def test_read_options(self, execute_kwargs): patch("pyathena.arrow.cursor.AthenaArrowResultSet") as result_set_class, ): cursor.execute("SELECT 1", **execute_kwargs) - first_kwargs = result_set_class.call_args.kwargs.copy() - cursor.execute("SELECT 1") - kwargs = first_kwargs + kwargs = result_set_class.call_args.kwargs expected = {**cursor_kwargs, **execute_kwargs} assert {key: kwargs[key] for key in expected} == expected - - defaults = cursor_kwargs.copy() - second_kwargs = result_set_class.call_args.kwargs - assert {key: second_kwargs[key] for key in defaults} == defaults diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index e342acf0d..90cf33c9a 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -4,27 +4,14 @@ # See LICENSE or https://opensource.org/licenses/MIT. # # SPDX-License-Identifier: MIT -import threading -from concurrent.futures import ThreadPoolExecutor -from datetime import UTC, datetime -from io import BytesIO +from datetime import datetime from unittest.mock import MagicMock, patch -import boto3 import pyarrow as pa -import pyarrow.parquet as pq -import pytest -from botocore.config import Config -from botocore.exceptions import ClientError -from botocore.response import StreamingBody -from fsspec.implementations.memory import MemoryFileSystem from pyathena.arrow.converter import DefaultArrowTypeConverter from pyathena.arrow.result_set import AthenaArrowResultSet -from pyathena.connection import Connection -from pyathena.filesystem.s3_executor import S3ThreadPoolExecutor from pyathena.model import AthenaQueryExecution -from pyathena.result_set import AthenaResultSet from pyathena.util import RetryConfig @@ -67,183 +54,3 @@ def test_read_csv_timestamps_with_the_same_name(self): ("a", datetime(2020, 1, 2, 3, 4, 5, 123456), datetime(2020, 1, 2, 3, 4, 5, 100000)), ("b", None, None), ] - - @pytest.mark.parametrize("unload", [False, True]) - def test_s3_workers_read_files(self, unload): - """Read real CSV/Parquet bytes through the S3 file and Arrow adapter offline.""" - if unload: - output = pa.BufferOutputStream() - pq.write_table(pa.table({"value": [1]}), output) - data = output.getvalue().to_pybytes() - key = "unload/part.parquet" - else: - data = b'"value"\n1\n' - key = "result.csv" - model_client = boto3.client( - "s3", region_name="us-east-1", aws_access_key_id="dummy", aws_secret_access_key="dummy" - ) - client = MagicMock() - client.meta.service_model = model_client.meta.service_model - client.meta.method_to_api_mapping = model_client.meta.method_to_api_mapping - model_client.close() - - def head_object(**kwargs): - if kwargs["Key"] != key: - raise ClientError( - {"Error": {"Code": "404"}, "ResponseMetadata": {"HTTPStatusCode": 404}}, - "HeadObject", - ) - return {"ContentLength": len(data), "LastModified": datetime(2026, 1, 1, tzinfo=UTC)} - - def get_object(**kwargs): - start, end = map(int, kwargs["Range"].removeprefix("bytes=").split("-")) - body = data[start : end + 1] - return {"Body": StreamingBody(BytesIO(body), len(body))} - - client.head_object.side_effect = head_object - client.get_object.side_effect = get_object - client.list_objects_v2.return_value = { - "Contents": [{"Key": key, "Size": len(data)}], - "IsTruncated": False, - } - connection = MagicMock(s3_client=client, retry_config=RetryConfig(attempt=1)) - execution = MagicMock( - state=AthenaQueryExecution.STATE_FAILED, - output_location=f"s3://bucket/{key}", - substatement_type=None, - ) - result = AthenaArrowResultSet( - connection=connection, - converter=DefaultArrowTypeConverter(), - query_execution=execution, - arraysize=1, - retry_config=RetryConfig(), - s3_max_workers=2, - unload=unload, - unload_location="s3://bucket/unload/", - ) - execution.state = AthenaQueryExecution.STATE_SUCCEEDED - result._metadata = ( - {"Name": "value", "Type": "integer", "Precision": 0, "Scale": 0, "Nullable": "UNKNOWN"}, - ) - with ( - patch.object(result, "_read_data_manifest", return_value=[f"s3://bucket/{key}"]), - patch( - "pyathena.filesystem.s3.S3ThreadPoolExecutor", wraps=S3ThreadPoolExecutor - ) as executor, - ): - table = result._as_arrow() - assert table.to_pylist() == [{"value": 1}] - assert executor.call_count > 0 - assert all(call.kwargs["max_workers"] == 2 for call in executor.call_args_list) - assert result._fs.handler.fs.max_workers == 2 - client.get_object.assert_called() - client.close.assert_not_called() - - @pytest.mark.parametrize("failure", [None, "filesystem", "read"]) - def test_s3_workers_owned_client_cleanup(self, failure): - """Dedicated timeout clients close after eager reads, including failures.""" - connection = MagicMock(region_name="us-east-1") - connection.s3_config = Config(max_pool_connections=17, read_timeout=60) - connection._s3_client_kwargs = {"use_ssl": True} - owned_client = connection.session.client.return_value - filesystem = MemoryFileSystem(skip_instance_cache=True) - execution = MagicMock( - state=AthenaQueryExecution.STATE_SUCCEEDED, output_location="s3://bucket/result.csv" - ) - with ( - patch.object(AthenaResultSet, "_pre_fetch"), - patch("pyathena.arrow.result_set.S3FileSystem", return_value=filesystem) as factory, - patch.object( - AthenaArrowResultSet, "_as_arrow", return_value=pa.table({"value": [1]}) - ) as read, - ): - if failure == "filesystem": - factory.side_effect = RuntimeError("filesystem failed") - elif failure == "read": - read.side_effect = RuntimeError("read failed") - kwargs = { - "connection": connection, - "converter": DefaultArrowTypeConverter(), - "query_execution": execution, - "arraysize": 1, - "retry_config": RetryConfig(), - "s3_max_workers": 2, - "connect_timeout": 1.5, - "request_timeout": 4.5, - } - if failure: - with pytest.raises(RuntimeError, match="failed"): - AthenaArrowResultSet(**kwargs) - else: - result = AthenaArrowResultSet(**kwargs) - assert result.as_arrow().to_pylist() == [{"value": 1}] - config = connection.session.client.call_args.kwargs["config"] - assert config.connect_timeout == 1.5 - assert config.read_timeout == 4.5 - assert config.max_pool_connections == 17 - assert connection.s3_config.read_timeout == 60 - assert connection.session.client.call_args.kwargs["use_ssl"] is True - assert factory.call_args.kwargs["s3_client"] is owned_client - owned_client.close.assert_called_once_with() - connection.s3_client.close.assert_not_called() - - def test_s3_workers_native_default(self): - """The default still builds the native Arrow filesystem, without a boto3 client.""" - connection = MagicMock(region_name="us-east-1", profile_name=None) - connection._kwargs = {"aws_access_key_id": "dummy", "aws_secret_access_key": "dummy"} - with patch("pyarrow.fs.S3FileSystem") as native: - AthenaArrowResultSet( - connection=connection, - converter=DefaultArrowTypeConverter(), - query_execution=MagicMock(state=AthenaQueryExecution.STATE_FAILED), - arraysize=1, - retry_config=RetryConfig(), - ) - native.assert_called_once_with( - access_key="dummy", secret_key="dummy", session_token=None, region="us-east-1" - ) - connection.session.client.assert_not_called() - - def test_s3_workers_serializes_client_creation(self): - """Dedicated and shared S3 clients use the same creation lock.""" - session = MagicMock() - connection = Connection( - session=session, region_name="us-east-1", s3_staging_dir="s3://bucket/path/" - ) - created = [] - - def create_client(*args, **kwargs): - assert args == ("s3",) - assert connection._s3_client_lock.locked() - client = MagicMock() - created.append(client) - return client - - session.client.side_effect = create_client - barrier = threading.Barrier(4) - - def create_result(index): - barrier.wait(timeout=5) - if index == 0: - return connection.s3_client - return AthenaArrowResultSet( - connection=connection, - converter=DefaultArrowTypeConverter(), - query_execution=MagicMock(state=AthenaQueryExecution.STATE_FAILED), - arraysize=1, - retry_config=RetryConfig(), - s3_max_workers=2, - request_timeout=4.5, - ) - - with ThreadPoolExecutor(max_workers=4) as executor: - results = list(executor.map(create_result, range(4))) - assert len(created) == 4 - shared_client = results[0] - for client in created: - if client is shared_client: - client.close.assert_not_called() - else: - client.close.assert_called_once_with() - assert not connection._s3_client_lock.locked() diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 9051fc5b2..410c6f697 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -115,20 +115,6 @@ def __getitem__(self, item): return BarrierDirCache() - @pytest.mark.parametrize("with_connection", [False, True]) - def test_s3_workers_client_override(self, with_connection): - """An injected client bypasses client creation and preserves connection retries.""" - connection = mock.MagicMock() if with_connection else None - client = mock.MagicMock() - with mock.patch.object(S3FileSystem, "_get_client_compatible_with_s3fs") as factory: - fs = S3FileSystem(connection=connection, s3_client=client, skip_instance_cache=True) - factory.assert_not_called() - assert fs.core.client is client - if connection is not None: - assert fs.core.retry_config is connection.retry_config - assert "s3_client" not in vars(connection) - client.close.assert_not_called() - def test_get_client_compatible_with_s3fs(self): # Only constructs a boto3 client; no AWS access. fs = S3FileSystem( diff --git a/tests/pyathena/pandas/test_async_cursor.py b/tests/pyathena/pandas/test_async_cursor.py index fce96b16f..86fa37146 100644 --- a/tests/pyathena/pandas/test_async_cursor.py +++ b/tests/pyathena/pandas/test_async_cursor.py @@ -669,18 +669,10 @@ def test_read_options(self, execute_kwargs): ): _, future = cursor.execute("SELECT 1", **execute_kwargs) future.result() - first_kwargs = result_set_class.call_args.kwargs.copy() - _, future = cursor.execute("SELECT 1") - future.result() - kwargs = first_kwargs + kwargs = result_set_class.call_args.kwargs expected = {**cursor_kwargs, **execute_kwargs} expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected - - defaults = cursor_kwargs.copy() - defaults["max_workers"] = defaults.pop("s3_max_workers") - second_kwargs = result_set_class.call_args.kwargs - assert {key: second_kwargs[key] for key in defaults} == defaults assert cursor._executor._max_workers == 1 @pytest.mark.parametrize( diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index 1d353dd49..912ee6158 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -1938,18 +1938,11 @@ def test_read_options(self, execute_kwargs): patch("pyathena.pandas.cursor.AthenaPandasResultSet") as result_set_class, ): cursor.execute("SELECT 1", **execute_kwargs) - first_kwargs = result_set_class.call_args.kwargs.copy() - cursor.execute("SELECT 1") - kwargs = first_kwargs + kwargs = result_set_class.call_args.kwargs expected = {**cursor_kwargs, **execute_kwargs} expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected - defaults = cursor_kwargs.copy() - defaults["max_workers"] = defaults.pop("s3_max_workers") - second_kwargs = result_set_class.call_args.kwargs - assert {key: second_kwargs[key] for key in defaults} == defaults - @pytest.mark.parametrize( "pandas_cursor", [ diff --git a/tests/pyathena/polars/test_async_cursor.py b/tests/pyathena/polars/test_async_cursor.py index 9400eadf9..68290a3a1 100644 --- a/tests/pyathena/polars/test_async_cursor.py +++ b/tests/pyathena/polars/test_async_cursor.py @@ -357,16 +357,8 @@ def test_read_options(self, execute_kwargs): ): _, future = cursor.execute("SELECT 1", **execute_kwargs) future.result() - first_kwargs = result_set_class.call_args.kwargs.copy() - _, future = cursor.execute("SELECT 1") - future.result() - kwargs = first_kwargs + kwargs = result_set_class.call_args.kwargs expected = {**cursor_kwargs, **execute_kwargs} expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected - - defaults = cursor_kwargs.copy() - defaults["max_workers"] = defaults.pop("s3_max_workers") - second_kwargs = result_set_class.call_args.kwargs - assert {key: second_kwargs[key] for key in defaults} == defaults assert cursor._executor._max_workers == 1 diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 29d36948c..92df25567 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -908,14 +908,7 @@ def test_read_options(self, execute_kwargs): patch("pyathena.polars.cursor.AthenaPolarsResultSet") as result_set_class, ): cursor.execute("SELECT 1", **execute_kwargs) - first_kwargs = result_set_class.call_args.kwargs.copy() - cursor.execute("SELECT 1") - kwargs = first_kwargs + kwargs = result_set_class.call_args.kwargs expected = {**cursor_kwargs, **execute_kwargs} expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected - - defaults = cursor_kwargs.copy() - defaults["max_workers"] = defaults.pop("s3_max_workers") - second_kwargs = result_set_class.call_args.kwargs - assert {key: second_kwargs[key] for key in defaults} == defaults diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index 187dfa272..1c5eed6fa 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -1,21 +1,10 @@ -from inspect import isawaitable from typing import Any -from unittest.mock import MagicMock, patch import pytest import tenacity.nap from botocore.exceptions import ClientError from pyathena import DataError -from pyathena.aio.arrow.cursor import AioArrowCursor -from pyathena.aio.pandas.cursor import AioPandasCursor -from pyathena.aio.polars.cursor import AioPolarsCursor -from pyathena.arrow.async_cursor import AsyncArrowCursor -from pyathena.arrow.cursor import ArrowCursor -from pyathena.pandas.async_cursor import AsyncPandasCursor -from pyathena.pandas.cursor import PandasCursor -from pyathena.polars.async_cursor import AsyncPolarsCursor -from pyathena.polars.cursor import PolarsCursor from pyathena.util import ( RetryConfig, _is_throttling_error, @@ -383,76 +372,3 @@ def call(): assert caught.value is error assert calls == expected_calls - - -S3_WORKER_CURSORS = ( - PandasCursor, - AsyncPandasCursor, - AioPandasCursor, - PolarsCursor, - AsyncPolarsCursor, - AioPolarsCursor, - ArrowCursor, - AsyncArrowCursor, - AioArrowCursor, -) - - -async def _execute_s3_worker_query(cursor, kwargs): - result = cursor.execute("SELECT 1", **kwargs) - if isawaitable(result): - await result - - -def _s3_worker_cursor(cursor_class, **kwargs): - return cursor_class( - connection=MagicMock(), - converter=MagicMock(), - formatter=MagicMock(), - retry_config=RetryConfig(), - **kwargs, - ) - - -@pytest.mark.parametrize("cursor_class", S3_WORKER_CURSORS) -@pytest.mark.parametrize( - ("value", "error"), [(0, ValueError), (-1, ValueError), (True, TypeError), (1.5, TypeError)] -) -def test_s3_workers_invalid_constructor(cursor_class, value, error): - with pytest.raises(error, match="s3_max_workers"): - _s3_worker_cursor(cursor_class, s3_max_workers=value) - - -@pytest.mark.parametrize("cursor_class", S3_WORKER_CURSORS) -@pytest.mark.parametrize( - "kwargs", [{"s3_max_workers": 0}, {"s3_max_workers": -1}, {"max_workers": 2}] -) -async def test_s3_workers_invalid_execute(cursor_class, kwargs): - cursor = _s3_worker_cursor(cursor_class) - error = TypeError if "max_workers" in kwargs else ValueError - try: - with patch.object(cursor, "_execute") as execute: - with pytest.raises(error, match="s3_max_workers"): - await _execute_s3_worker_query(cursor, kwargs) - execute.assert_not_called() - finally: - closed = cursor.close() - if isawaitable(closed): - await closed - - -@pytest.mark.parametrize( - "cursor_class", - [PandasCursor, PolarsCursor, ArrowCursor, AioPandasCursor, AioPolarsCursor, AioArrowCursor], -) -def test_s3_workers_reject_former_constructor_keyword(cursor_class): - with pytest.raises(TypeError, match="s3_max_workers"): - _s3_worker_cursor(cursor_class, max_workers=2) - - -@pytest.mark.parametrize( - "cursor_class", [PandasCursor, PolarsCursor, AioPandasCursor, AioPolarsCursor] -) -def test_s3_workers_none_is_arrow_only(cursor_class): - with pytest.raises(TypeError, match="s3_max_workers"): - _s3_worker_cursor(cursor_class, s3_max_workers=None) From f316843f84e1dc839e1dd243e1ca3fa76b188c18 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 01:13:19 +0900 Subject: [PATCH 5/6] Note that chunked Polars reads do not use s3_max_workers Co-Authored-By: Claude Opus 5.5 --- docs/usage.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/usage.md b/docs/usage.md index d54bb3a12..ef3a28323 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -45,7 +45,8 @@ The cursor class can still be selected positionally with `connection.cursor(Dict The pandas and Polars cursors take the number of threads that read each result file from S3 as `s3_max_workers`. It can be set in the cursor constructor and overridden for one query in `execute()`. The default is `(cpu_count() or 1) * 5`. -Polars reads UNLOAD results with its native Parquet reader, which does not use this setting. +Polars uses it only for CSV results read without `chunksize`. +Chunked CSV reads and UNLOAD results use Polars' native readers, which do not use this setting. The thread-pool cursors, such as `AsyncCursor`, `AsyncPandasCursor`, and `AsyncPolarsCursor`, keep `max_workers` for the size of the thread pool that waits for queries and collects their results. From 9df509eaa4a8432db13e62837d4613948352be76 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 12:17:21 +0900 Subject: [PATCH 6/6] Describe where the former max_workers keyword fails Co-Authored-By: Claude Opus 5.5 --- docs/usage.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/usage.md b/docs/usage.md index 02aa9d8e8..61a6fa0c8 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -65,7 +65,8 @@ query_id, future = cursor.execute("SELECT 1", s3_max_workers=3) Before PyAthena 4.0, `max_workers` set the S3 read workers of `PandasCursor`, `PolarsCursor`, `AioPandasCursor`, and `AioPolarsCursor`, and of `execute()` in the pandas and Polars cursors. Replace it with `s3_max_workers` in those places. -In these cursors, `max_workers` now raises `TypeError` in the constructor, and in `execute()` after the query has run. +Passing `max_workers` to those four constructors raises `TypeError`. +Passing it to `execute()` raises `TypeError` after the query has run; with `AsyncPandasCursor` and `AsyncPolarsCursor`, the returned future raises it. `AsyncPolarsCursor` used its `max_workers` for both the thread pool and the S3 read workers; pass both arguments to keep the former values. ## Managed query result storage