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..5eb67dd7a 100644 --- a/benchmarks/tests/test_adapters.py +++ b/benchmarks/tests/test_adapters.py @@ -129,13 +129,18 @@ 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 cursor.close() @pytest.mark.parametrize("api", ["sync", "thread", "aio"]) diff --git a/docs/usage.md b/docs/usage.md index 3f6be5a25..61a6fa0c8 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -46,6 +46,29 @@ A cursor subclass that receives settings of its own must not pass them on to the The pandas and Polars cursors pass the `execute()` keyword arguments that they do not use themselves to the reader, such as `pandas.read_csv()`. The reader rejects an unknown name when it reads the query results, which PyAthena raises as `OperationalError`; the name is not checked when there are no results to read. +## 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 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. + +```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. +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 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/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index 4000d4baa..cc7f259d3 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -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 @@ -110,7 +110,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 +168,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 @@ -218,7 +218,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..17a551a99 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -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,7 +78,7 @@ 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``, @@ -100,7 +100,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 +152,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`` @@ -197,7 +197,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 7d56e64e1..5c06ef97b 100644 --- a/pyathena/arrow/async_cursor.py +++ b/pyathena/arrow/async_cursor.py @@ -93,7 +93,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. diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index 5f047933c..2bd74f5d7 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/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index 01c968052..4337065d5 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -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,6 +111,7 @@ 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__``. """ @@ -128,6 +130,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 +187,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,7 +241,7 @@ 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. diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 1ed437935..37b68c044 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -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 @@ -133,7 +133,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 +194,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 @@ -247,7 +247,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..8f015621e 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -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: @@ -123,6 +125,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 +189,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 +234,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`` diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index a7e31a9f2..692c31d2a 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -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(). @@ -128,7 +128,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 +190,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`` @@ -237,7 +237,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/tests/pyathena/aio/pandas/test_cursor.py b/tests/pyathena/aio/pandas/test_cursor.py index ecb123a34..5fd863a62 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( @@ -179,4 +179,5 @@ async def test_read_options(self, execute_kwargs): await cursor.execute("SELECT 1", **execute_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 diff --git a/tests/pyathena/aio/polars/test_cursor.py b/tests/pyathena/aio/polars/test_cursor.py index 051d078dc..9c9003c74 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( @@ -165,4 +165,5 @@ async def test_read_options(self, execute_kwargs): await cursor.execute("SELECT 1", **execute_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 diff --git a/tests/pyathena/pandas/test_async_cursor.py b/tests/pyathena/pandas/test_async_cursor.py index 590d9d8b8..86fa37146 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"), @@ -665,7 +671,9 @@ def test_read_options(self, execute_kwargs): future.result() 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 + assert cursor._executor._max_workers == 1 @pytest.mark.parametrize( "async_pandas_cursor", diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index ad7bce11d..912ee6158 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( @@ -1940,6 +1940,7 @@ def test_read_options(self, execute_kwargs): cursor.execute("SELECT 1", **execute_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 @pytest.mark.parametrize( diff --git a/tests/pyathena/polars/test_async_cursor.py b/tests/pyathena/polars/test_async_cursor.py index ac4a761d9..68290a3a1 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"), @@ -358,4 +359,6 @@ def test_read_options(self, execute_kwargs): future.result() 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 + assert cursor._executor._max_workers == 1 diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index d739bc22f..92df25567 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( @@ -910,4 +910,5 @@ def test_read_options(self, execute_kwargs): cursor.execute("SELECT 1", **execute_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