PyAthena 3.29.1__tar.gz → 3.29.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (87) hide show
  1. {pyathena-3.29.1 → pyathena-3.29.3}/PKG-INFO +3 -3
  2. {pyathena-3.29.1 → pyathena-3.29.3}/README.md +2 -2
  3. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/__init__.py +2 -2
  4. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/_version.py +2 -2
  5. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/arrow/cursor.py +1 -1
  6. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/pandas/cursor.py +1 -1
  7. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/polars/cursor.py +1 -1
  8. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/s3fs/cursor.py +8 -5
  9. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/spark/cursor.py +1 -1
  10. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/filesystem/s3.py +31 -22
  11. pyathena-3.29.3/pyathena/filesystem/s3_async.py +357 -0
  12. pyathena-3.29.3/pyathena/filesystem/s3_executor.py +89 -0
  13. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/result_set.py +6 -2
  14. {pyathena-3.29.1 → pyathena-3.29.3}/pyproject.toml +10 -4
  15. {pyathena-3.29.1 → pyathena-3.29.3}/.gitignore +0 -0
  16. {pyathena-3.29.1 → pyathena-3.29.3}/LICENSE +0 -0
  17. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/__init__.py +0 -0
  18. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/arrow/__init__.py +0 -0
  19. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/common.py +0 -0
  20. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/connection.py +0 -0
  21. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/cursor.py +0 -0
  22. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/pandas/__init__.py +0 -0
  23. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/polars/__init__.py +0 -0
  24. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/result_set.py +0 -0
  25. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/s3fs/__init__.py +0 -0
  26. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/spark/__init__.py +0 -0
  27. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/__init__.py +0 -0
  28. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/arrow.py +0 -0
  29. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/base.py +0 -0
  30. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/pandas.py +0 -0
  31. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/polars.py +0 -0
  32. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/rest.py +0 -0
  33. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/sqlalchemy/s3fs.py +0 -0
  34. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/aio/util.py +0 -0
  35. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/__init__.py +0 -0
  36. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/async_cursor.py +0 -0
  37. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/converter.py +0 -0
  38. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/cursor.py +0 -0
  39. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/result_set.py +0 -0
  40. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/arrow/util.py +0 -0
  41. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/async_cursor.py +0 -0
  42. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/common.py +0 -0
  43. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/connection.py +0 -0
  44. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/converter.py +0 -0
  45. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/cursor.py +0 -0
  46. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/error.py +0 -0
  47. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/filesystem/__init__.py +0 -0
  48. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/filesystem/s3_object.py +0 -0
  49. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/formatter.py +0 -0
  50. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/model.py +0 -0
  51. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/__init__.py +0 -0
  52. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/async_cursor.py +0 -0
  53. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/converter.py +0 -0
  54. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/cursor.py +0 -0
  55. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/result_set.py +0 -0
  56. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/pandas/util.py +0 -0
  57. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/__init__.py +0 -0
  58. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/async_cursor.py +0 -0
  59. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/converter.py +0 -0
  60. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/cursor.py +0 -0
  61. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/result_set.py +0 -0
  62. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/polars/util.py +0 -0
  63. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/py.typed +0 -0
  64. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/result_set.py +0 -0
  65. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/__init__.py +0 -0
  66. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/async_cursor.py +0 -0
  67. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/converter.py +0 -0
  68. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/cursor.py +0 -0
  69. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/s3fs/reader.py +0 -0
  70. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/spark/__init__.py +0 -0
  71. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/spark/async_cursor.py +0 -0
  72. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/spark/common.py +0 -0
  73. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/spark/cursor.py +0 -0
  74. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/__init__.py +0 -0
  75. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/arrow.py +0 -0
  76. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/base.py +0 -0
  77. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/compiler.py +0 -0
  78. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/constants.py +0 -0
  79. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/pandas.py +0 -0
  80. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/polars.py +0 -0
  81. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/preparer.py +0 -0
  82. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/requirements.py +0 -0
  83. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/rest.py +0 -0
  84. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/s3fs.py +0 -0
  85. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/types.py +0 -0
  86. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/sqlalchemy/util.py +0 -0
  87. {pyathena-3.29.1 → pyathena-3.29.3}/pyathena/util.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: PyAthena
3
- Version: 3.29.1
3
+ Version: 3.29.3
4
4
  Summary: Python DB API 2.0 (PEP 249) client for Amazon Athena
5
5
  Project-URL: homepage, https://github.com/pyathena-dev/PyAthena/
6
6
  Project-URL: repository, https://github.com/pyathena-dev/PyAthena/
@@ -115,10 +115,10 @@ Native asyncio is also supported:
115
115
 
116
116
  ```python
117
117
  import asyncio
118
- from pyathena import aconnect
118
+ from pyathena import aio_connect
119
119
 
120
120
  async def main():
121
- async with await aconnect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
121
+ async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
122
122
  region_name="us-west-2") as conn:
123
123
  cursor = conn.cursor()
124
124
  await cursor.execute("SELECT 1")
@@ -56,10 +56,10 @@ Native asyncio is also supported:
56
56
 
57
57
  ```python
58
58
  import asyncio
59
- from pyathena import aconnect
59
+ from pyathena import aio_connect
60
60
 
61
61
  async def main():
62
- async with await aconnect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
62
+ async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
63
63
  region_name="us-west-2") as conn:
64
64
  cursor = conn.cursor()
65
65
  await cursor.execute("SELECT 1")
@@ -131,7 +131,7 @@ def connect(*args, **kwargs) -> "Connection[Any]":
131
131
  return Connection(*args, **kwargs)
132
132
 
133
133
 
134
- async def aconnect(*args, **kwargs) -> "AioConnection":
134
+ async def aio_connect(*args, **kwargs) -> "AioConnection":
135
135
  """Create a new async database connection to Amazon Athena.
136
136
 
137
137
  This is the async counterpart of :func:`connect`. It returns an
@@ -147,7 +147,7 @@ async def aconnect(*args, **kwargs) -> "AioConnection":
147
147
 
148
148
  Example:
149
149
  >>> import pyathena
150
- >>> conn = await pyathena.aconnect(
150
+ >>> conn = await pyathena.aio_connect(
151
151
  ... s3_staging_dir='s3://my-bucket/staging/',
152
152
  ... region_name='us-east-1',
153
153
  ... )
@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
28
28
  commit_id: COMMIT_ID
29
29
  __commit_id__: COMMIT_ID
30
30
 
31
- __version__ = version = '3.29.1'
32
- __version_tuple__ = version_tuple = (3, 29, 1)
31
+ __version__ = version = '3.29.3'
32
+ __version_tuple__ = version_tuple = (3, 29, 3)
33
33
 
34
34
  __commit_id__ = commit_id = None
@@ -29,7 +29,7 @@ class AioArrowCursor(WithAsyncFetch):
29
29
  operations, keeping the event loop free.
30
30
 
31
31
  Example:
32
- >>> async with await pyathena.aconnect(...) as conn:
32
+ >>> async with await pyathena.aio_connect(...) as conn:
33
33
  ... cursor = conn.cursor(AioArrowCursor)
34
34
  ... await cursor.execute("SELECT * FROM my_table")
35
35
  ... table = cursor.as_arrow()
@@ -40,7 +40,7 @@ class AioPandasCursor(WithAsyncFetch):
40
40
  when ``chunksize`` is set, as fetch calls trigger lazy S3 reads.
41
41
 
42
42
  Example:
43
- >>> async with await pyathena.aconnect(...) as conn:
43
+ >>> async with await pyathena.aio_connect(...) as conn:
44
44
  ... cursor = conn.cursor(AioPandasCursor)
45
45
  ... await cursor.execute("SELECT * FROM my_table")
46
46
  ... df = cursor.as_pandas()
@@ -31,7 +31,7 @@ class AioPolarsCursor(WithAsyncFetch):
31
31
  when ``chunksize`` is set, as fetch calls trigger lazy S3 reads.
32
32
 
33
33
  Example:
34
- >>> async with await pyathena.aconnect(...) as conn:
34
+ >>> async with await pyathena.aio_connect(...) as conn:
35
35
  ... cursor = conn.cursor(AioPolarsCursor)
36
36
  ... await cursor.execute("SELECT * FROM my_table")
37
37
  ... df = cursor.as_polars()
@@ -8,6 +8,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union, cast
8
8
  from pyathena.aio.common import WithAsyncFetch
9
9
  from pyathena.common import CursorIterator
10
10
  from pyathena.error import OperationalError, ProgrammingError
11
+ from pyathena.filesystem.s3_async import AioS3FileSystem
11
12
  from pyathena.model import AthenaQueryExecution
12
13
  from pyathena.s3fs.converter import DefaultS3FSTypeConverter
13
14
  from pyathena.s3fs.result_set import AthenaS3FSResultSet, CSVReaderType
@@ -16,14 +17,15 @@ _logger = logging.getLogger(__name__) # type: ignore
16
17
 
17
18
 
18
19
  class AioS3FSCursor(WithAsyncFetch):
19
- """Native asyncio cursor that reads CSV results via S3FileSystem.
20
+ """Native asyncio cursor that reads CSV results via AioS3FileSystem.
20
21
 
21
- Uses ``asyncio.to_thread()`` for result set creation and fetch operations
22
- because ``AthenaS3FSResultSet`` lazily streams rows from S3 via a CSV
23
- reader, making fetch calls blocking I/O.
22
+ Uses ``AioS3FileSystem`` for S3 operations, which replaces
23
+ ``ThreadPoolExecutor`` parallelism with ``asyncio.gather`` +
24
+ ``asyncio.to_thread``. Fetch operations are wrapped in
25
+ ``asyncio.to_thread()`` because CSV reading is blocking I/O.
24
26
 
25
27
  Example:
26
- >>> async with await pyathena.aconnect(...) as conn:
28
+ >>> async with await pyathena.aio_connect(...) as conn:
27
29
  ... cursor = conn.cursor(AioS3FSCursor)
28
30
  ... await cursor.execute("SELECT * FROM my_table")
29
31
  ... row = await cursor.fetchone()
@@ -127,6 +129,7 @@ class AioS3FSCursor(WithAsyncFetch):
127
129
  arraysize=self.arraysize,
128
130
  retry_config=self._retry_config,
129
131
  csv_reader=self._csv_reader,
132
+ filesystem_class=AioS3FileSystem,
130
133
  **kwargs,
131
134
  )
132
135
  else:
@@ -33,7 +33,7 @@ class AioSparkCursor(SparkBaseCursor, WithCalculationExecution):
33
33
 
34
34
  Example:
35
35
  >>> import asyncio
36
- >>> async with await pyathena.aconnect(
36
+ >>> async with await pyathena.aio_connect(
37
37
  ... work_group="spark-workgroup",
38
38
  ... cursor_class=AioSparkCursor,
39
39
  ... ) as conn:
@@ -1,13 +1,11 @@
1
1
  # -*- coding: utf-8 -*-
2
2
  from __future__ import annotations
3
3
 
4
- import itertools
5
4
  import logging
6
5
  import mimetypes
7
6
  import os.path
8
7
  import re
9
8
  from concurrent.futures import Future, as_completed
10
- from concurrent.futures.thread import ThreadPoolExecutor
11
9
  from copy import deepcopy
12
10
  from datetime import datetime
13
11
  from multiprocessing import cpu_count
@@ -23,6 +21,7 @@ from fsspec.spec import AbstractBufferedFile
23
21
  from fsspec.utils import tokenize
24
22
 
25
23
  import pyathena
24
+ from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPoolExecutor
26
25
  from pyathena.filesystem.s3_object import (
27
26
  S3CompleteMultipartUpload,
28
27
  S3MultipartUpload,
@@ -686,6 +685,20 @@ class S3FileSystem(AbstractFileSystem):
686
685
  **request,
687
686
  )
688
687
 
688
+ def _create_executor(self, max_workers: int) -> S3Executor:
689
+ """Create an executor strategy for parallel operations.
690
+
691
+ Subclasses can override to provide alternative execution strategies
692
+ (e.g., asyncio-based execution).
693
+
694
+ Args:
695
+ max_workers: Maximum number of parallel workers.
696
+
697
+ Returns:
698
+ An S3Executor instance.
699
+ """
700
+ return S3ThreadPoolExecutor(max_workers=max_workers)
701
+
689
702
  def _delete_objects(
690
703
  self, bucket: str, paths: List[str], max_workers: Optional[int] = None, **kwargs
691
704
  ) -> None:
@@ -703,7 +716,7 @@ class S3FileSystem(AbstractFileSystem):
703
716
  object_.update({"VersionId": version_id})
704
717
  delete_objects.append(object_)
705
718
 
706
- with ThreadPoolExecutor(max_workers=max_workers) as executor:
719
+ with self._create_executor(max_workers=max_workers) as executor:
707
720
  fs = []
708
721
  for delete in [
709
722
  delete_objects[i : i + self.DELETE_OBJECTS_MAX_KEYS]
@@ -861,7 +874,7 @@ class S3FileSystem(AbstractFileSystem):
861
874
  **kwargs,
862
875
  )
863
876
  parts = []
864
- with ThreadPoolExecutor(max_workers=max_workers) as executor:
877
+ with self._create_executor(max_workers=max_workers) as executor:
865
878
  fs = [
866
879
  executor.submit(
867
880
  self._upload_part_copy,
@@ -1106,6 +1119,7 @@ class S3FileSystem(AbstractFileSystem):
1106
1119
  mode,
1107
1120
  version_id=None,
1108
1121
  max_workers=max_workers,
1122
+ executor=self._create_executor(max_workers=max_workers),
1109
1123
  block_size=block_size,
1110
1124
  cache_type=cache_type,
1111
1125
  autocommit=autocommit,
@@ -1256,6 +1270,7 @@ class S3File(AbstractBufferedFile):
1256
1270
  mode: str = "rb",
1257
1271
  version_id: Optional[str] = None,
1258
1272
  max_workers: int = (cpu_count() or 1) * 5,
1273
+ executor: Optional[S3Executor] = None,
1259
1274
  block_size: int = S3FileSystem.DEFAULT_BLOCK_SIZE,
1260
1275
  cache_type: str = "bytes",
1261
1276
  autocommit: bool = True,
@@ -1265,7 +1280,7 @@ class S3File(AbstractBufferedFile):
1265
1280
  **kwargs,
1266
1281
  ) -> None:
1267
1282
  self.max_workers = max_workers
1268
- self._executor = ThreadPoolExecutor(max_workers=max_workers)
1283
+ self._executor: S3Executor = executor or S3ThreadPoolExecutor(max_workers=max_workers)
1269
1284
  self.s3_additional_kwargs = s3_additional_kwargs if s3_additional_kwargs else {}
1270
1285
 
1271
1286
  super().__init__(
@@ -1481,24 +1496,18 @@ class S3File(AbstractBufferedFile):
1481
1496
  start, end, max_workers=self.max_workers, worker_block_size=self.blocksize
1482
1497
  )
1483
1498
  if len(ranges) > 1:
1484
- object_ = self._merge_objects(
1485
- list(
1486
- self._executor.map(
1487
- lambda bucket, key, ranges, version_id, kwargs: self.fs._get_object(
1488
- bucket=bucket,
1489
- key=key,
1490
- ranges=ranges,
1491
- version_id=version_id,
1492
- **kwargs,
1493
- ),
1494
- itertools.repeat(self.bucket),
1495
- itertools.repeat(self.key),
1496
- ranges,
1497
- itertools.repeat(self.version_id),
1498
- itertools.repeat(self.s3_additional_kwargs),
1499
- )
1499
+ futures = [
1500
+ self._executor.submit(
1501
+ self.fs._get_object,
1502
+ bucket=self.bucket,
1503
+ key=self.key,
1504
+ ranges=r,
1505
+ version_id=self.version_id,
1506
+ **self.s3_additional_kwargs,
1500
1507
  )
1501
- )
1508
+ for r in ranges
1509
+ ]
1510
+ object_ = self._merge_objects([f.result() for f in as_completed(futures)])
1502
1511
  else:
1503
1512
  object_ = self.fs._get_object(
1504
1513
  self.bucket,
@@ -0,0 +1,357 @@
1
+ # -*- coding: utf-8 -*-
2
+ from __future__ import annotations
3
+
4
+ import asyncio
5
+ import logging
6
+ from multiprocessing import cpu_count
7
+ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
8
+
9
+ from fsspec.asyn import AsyncFileSystem
10
+ from fsspec.callbacks import _DEFAULT_CALLBACK
11
+
12
+ from pyathena.filesystem.s3 import S3File, S3FileSystem
13
+ from pyathena.filesystem.s3_executor import S3AioExecutor
14
+ from pyathena.filesystem.s3_object import S3Object
15
+
16
+ if TYPE_CHECKING:
17
+ from datetime import datetime
18
+
19
+ from pyathena.connection import Connection
20
+
21
+ _logger = logging.getLogger(__name__)
22
+
23
+
24
+ class AioS3FileSystem(AsyncFileSystem):
25
+ """An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem.
26
+
27
+ This class wraps ``S3FileSystem`` to provide native asyncio support. Instead of
28
+ using ``ThreadPoolExecutor`` for parallel operations, it uses ``asyncio.gather``
29
+ with ``asyncio.to_thread`` for natural integration with the asyncio event loop.
30
+
31
+ The implementation uses composition: an internal ``S3FileSystem`` instance handles
32
+ all boto3 calls, while this class delegates to it via ``asyncio.to_thread()``.
33
+ This avoids diamond inheritance issues and keeps all boto3 logic in one place.
34
+
35
+ File handles created by ``_open`` use ``S3AioExecutor`` so that parallel
36
+ operations (range reads, multipart uploads) are dispatched via the event loop
37
+ instead of spawning additional threads.
38
+
39
+ Attributes:
40
+ _sync_fs: The internal synchronous S3FileSystem instance.
41
+
42
+ Example:
43
+ >>> from pyathena.filesystem.s3_async import AioS3FileSystem
44
+ >>> fs = AioS3FileSystem(asynchronous=True)
45
+ >>>
46
+ >>> # Use in async context
47
+ >>> files = await fs._ls('s3://my-bucket/data/')
48
+ >>>
49
+ >>> # Sync wrappers also available (auto-generated by fsspec)
50
+ >>> files = fs.ls('s3://my-bucket/data/')
51
+ """
52
+
53
+ # https://docs.aws.amazon.com/AmazonS3/latest/API/API_DeleteObjects.html
54
+ DELETE_OBJECTS_MAX_KEYS: int = 1000
55
+
56
+ protocol = ("s3", "s3a")
57
+ mirror_sync_methods = True
58
+ async_impl = True
59
+ _extra_tokenize_attributes = ("default_block_size",)
60
+
61
+ def __init__(
62
+ self,
63
+ connection: Optional["Connection[Any]"] = None,
64
+ default_block_size: Optional[int] = None,
65
+ default_cache_type: Optional[str] = None,
66
+ max_workers: int = (cpu_count() or 1) * 5,
67
+ s3_additional_kwargs: Optional[Dict[str, Any]] = None,
68
+ asynchronous: bool = False,
69
+ loop: Optional[Any] = None,
70
+ batch_size: Optional[int] = None,
71
+ **kwargs,
72
+ ) -> None:
73
+ super().__init__(
74
+ asynchronous=asynchronous,
75
+ loop=loop,
76
+ batch_size=batch_size,
77
+ **kwargs,
78
+ )
79
+ self._sync_fs = S3FileSystem(
80
+ connection=connection,
81
+ default_block_size=default_block_size,
82
+ default_cache_type=default_cache_type,
83
+ max_workers=max_workers,
84
+ s3_additional_kwargs=s3_additional_kwargs,
85
+ **kwargs,
86
+ )
87
+ # Share dircache for cache coherence between async and sync instances
88
+ self.dircache = self._sync_fs.dircache
89
+
90
+ @staticmethod
91
+ def parse_path(path: str) -> Tuple[str, Optional[str], Optional[str]]:
92
+ return S3FileSystem.parse_path(path)
93
+
94
+ async def _info(self, path: str, **kwargs) -> S3Object:
95
+ return await asyncio.to_thread(self._sync_fs.info, path, **kwargs)
96
+
97
+ async def _ls(
98
+ self, path: str, detail: bool = False, **kwargs
99
+ ) -> Union[List[S3Object], List[str]]:
100
+ return await asyncio.to_thread(self._sync_fs.ls, path, detail=detail, **kwargs)
101
+
102
+ async def _cat_file(
103
+ self, path: str, start: Optional[int] = None, end: Optional[int] = None, **kwargs
104
+ ) -> bytes:
105
+ return await asyncio.to_thread(self._sync_fs.cat_file, path, start=start, end=end, **kwargs)
106
+
107
+ async def _exists(self, path: str, **kwargs) -> bool:
108
+ return await asyncio.to_thread(self._sync_fs.exists, path, **kwargs)
109
+
110
+ async def _rm_file(self, path: str, **kwargs) -> None:
111
+ await asyncio.to_thread(self._sync_fs.rm_file, path, **kwargs)
112
+
113
+ async def _pipe_file(self, path: str, value: bytes, **kwargs) -> None:
114
+ await asyncio.to_thread(self._sync_fs.pipe_file, path, value, **kwargs)
115
+
116
+ async def _put_file(self, lpath: str, rpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
117
+ await asyncio.to_thread(self._sync_fs.put_file, lpath, rpath, callback=callback, **kwargs)
118
+
119
+ async def _get_file(self, rpath: str, lpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None:
120
+ await asyncio.to_thread(self._sync_fs.get_file, rpath, lpath, callback=callback, **kwargs)
121
+
122
+ async def _mkdir(self, path: str, create_parents: bool = True, **kwargs) -> None:
123
+ await asyncio.to_thread(self._sync_fs.mkdir, path, create_parents=create_parents, **kwargs)
124
+
125
+ async def _makedirs(self, path: str, exist_ok: bool = False) -> None:
126
+ await asyncio.to_thread(self._sync_fs.makedirs, path, exist_ok=exist_ok)
127
+
128
+ async def _rm(self, path: Union[str, List[str]], recursive: bool = False, **kwargs) -> None:
129
+ """Remove files or directories using async parallel batch deletion.
130
+
131
+ For multiple paths, chunks into batches of 1000 (S3 API limit) and uses
132
+ ``asyncio.gather`` with ``asyncio.to_thread`` instead of ThreadPoolExecutor.
133
+ """
134
+ if isinstance(path, str):
135
+ path = [path]
136
+
137
+ bucket, _, _ = self.parse_path(path[0])
138
+
139
+ expand_paths: List[str] = []
140
+ for p in path:
141
+ expanded = await asyncio.to_thread(self._sync_fs.expand_path, p, recursive=recursive)
142
+ expand_paths.extend(expanded)
143
+
144
+ if not expand_paths:
145
+ return
146
+
147
+ quiet = kwargs.pop("Quiet", True)
148
+ delete_objects: List[Dict[str, Any]] = []
149
+ for p in expand_paths:
150
+ _, key, version_id = self.parse_path(p)
151
+ if key:
152
+ object_: Dict[str, Any] = {"Key": key}
153
+ if version_id:
154
+ object_["VersionId"] = version_id
155
+ delete_objects.append(object_)
156
+
157
+ if not delete_objects:
158
+ return
159
+
160
+ chunks = [
161
+ delete_objects[i : i + self.DELETE_OBJECTS_MAX_KEYS]
162
+ for i in range(0, len(delete_objects), self.DELETE_OBJECTS_MAX_KEYS)
163
+ ]
164
+
165
+ async def _delete_chunk(chunk: List[Dict[str, Any]]) -> None:
166
+ request = {
167
+ "Bucket": bucket,
168
+ "Delete": {
169
+ "Objects": chunk,
170
+ "Quiet": quiet,
171
+ },
172
+ }
173
+ await asyncio.to_thread(
174
+ self._sync_fs._call, self._sync_fs._client.delete_objects, **request
175
+ )
176
+
177
+ await asyncio.gather(*[_delete_chunk(chunk) for chunk in chunks])
178
+
179
+ for p in expand_paths:
180
+ self._sync_fs.invalidate_cache(p)
181
+
182
+ async def _cp_file(self, path1: str, path2: str, **kwargs) -> None:
183
+ """Copy an S3 object, using async parallel multipart upload for large files."""
184
+ kwargs.pop("onerror", None)
185
+ bucket1, key1, version_id1 = self.parse_path(path1)
186
+ bucket2, key2, version_id2 = self.parse_path(path2)
187
+ if version_id2:
188
+ raise ValueError("Cannot copy to a versioned file.")
189
+ if not key1 or not key2:
190
+ raise ValueError("Cannot copy buckets.")
191
+
192
+ info1 = await self._info(path1)
193
+ size1 = info1.get("size", 0)
194
+ if size1 <= S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE:
195
+ await asyncio.to_thread(
196
+ self._sync_fs._copy_object,
197
+ bucket1=bucket1,
198
+ key1=key1,
199
+ version_id1=version_id1,
200
+ bucket2=bucket2,
201
+ key2=key2,
202
+ **kwargs,
203
+ )
204
+ else:
205
+ await self._copy_object_with_multipart_upload(
206
+ bucket1=bucket1,
207
+ key1=key1,
208
+ version_id1=version_id1,
209
+ size1=size1,
210
+ bucket2=bucket2,
211
+ key2=key2,
212
+ **kwargs,
213
+ )
214
+ self._sync_fs.invalidate_cache(path2)
215
+
216
+ async def _copy_object_with_multipart_upload(
217
+ self,
218
+ bucket1: str,
219
+ key1: str,
220
+ size1: int,
221
+ bucket2: str,
222
+ key2: str,
223
+ block_size: Optional[int] = None,
224
+ version_id1: Optional[str] = None,
225
+ **kwargs,
226
+ ) -> None:
227
+ block_size = block_size if block_size else S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE
228
+ if (
229
+ block_size < S3FileSystem.MULTIPART_UPLOAD_MIN_PART_SIZE
230
+ or block_size > S3FileSystem.MULTIPART_UPLOAD_MAX_PART_SIZE
231
+ ):
232
+ raise ValueError("Block size must be greater than 5MiB and less than 5GiB.")
233
+
234
+ copy_source: Dict[str, Any] = {
235
+ "Bucket": bucket1,
236
+ "Key": key1,
237
+ }
238
+ if version_id1:
239
+ copy_source["VersionId"] = version_id1
240
+
241
+ ranges = S3File._get_ranges(
242
+ 0,
243
+ size1,
244
+ self._sync_fs.max_workers,
245
+ block_size,
246
+ )
247
+ multipart_upload = await asyncio.to_thread(
248
+ self._sync_fs._create_multipart_upload,
249
+ bucket=bucket2,
250
+ key=key2,
251
+ **kwargs,
252
+ )
253
+
254
+ async def _upload_part(i: int, range_: Tuple[int, int]) -> Dict[str, Any]:
255
+ result = await asyncio.to_thread(
256
+ self._sync_fs._upload_part_copy,
257
+ bucket=bucket2,
258
+ key=key2,
259
+ copy_source=copy_source,
260
+ upload_id=cast(str, multipart_upload.upload_id),
261
+ part_number=i + 1,
262
+ copy_source_ranges=range_,
263
+ )
264
+ return {
265
+ "ETag": result.etag,
266
+ "PartNumber": result.part_number,
267
+ }
268
+
269
+ parts = await asyncio.gather(*[_upload_part(i, r) for i, r in enumerate(ranges)])
270
+ parts_list = sorted(parts, key=lambda x: x["PartNumber"])
271
+
272
+ await asyncio.to_thread(
273
+ self._sync_fs._complete_multipart_upload,
274
+ bucket=bucket2,
275
+ key=key2,
276
+ upload_id=cast(str, multipart_upload.upload_id),
277
+ parts=parts_list,
278
+ )
279
+
280
+ async def _find(
281
+ self,
282
+ path: str,
283
+ maxdepth: Optional[int] = None,
284
+ withdirs: bool = False,
285
+ **kwargs,
286
+ ) -> Union[Dict[str, S3Object], List[str]]:
287
+ detail = kwargs.pop("detail", False)
288
+ files = await asyncio.to_thread(
289
+ self._sync_fs._find, path, maxdepth=maxdepth, withdirs=withdirs, **kwargs
290
+ )
291
+ if detail:
292
+ return {f.name: f for f in files}
293
+ return [f.name for f in files]
294
+
295
+ def _open(
296
+ self,
297
+ path: str,
298
+ mode: str = "rb",
299
+ block_size: Optional[int] = None,
300
+ cache_type: Optional[str] = None,
301
+ autocommit: bool = True,
302
+ cache_options: Optional[Dict[Any, Any]] = None,
303
+ **kwargs,
304
+ ) -> "AioS3File":
305
+ if block_size is None:
306
+ block_size = self._sync_fs.default_block_size
307
+ if cache_type is None:
308
+ cache_type = self._sync_fs.default_cache_type
309
+ max_workers = kwargs.pop("max_worker", self._sync_fs.max_workers)
310
+ s3_additional_kwargs = kwargs.pop("s3_additional_kwargs", {})
311
+ s3_additional_kwargs.update(self._sync_fs.s3_additional_kwargs)
312
+
313
+ return AioS3File(
314
+ self._sync_fs,
315
+ path,
316
+ mode,
317
+ version_id=None,
318
+ max_workers=max_workers,
319
+ executor=S3AioExecutor(loop=self._loop),
320
+ block_size=block_size,
321
+ cache_type=cache_type,
322
+ autocommit=autocommit,
323
+ cache_options=cache_options,
324
+ s3_additional_kwargs=s3_additional_kwargs,
325
+ **kwargs,
326
+ )
327
+
328
+ def sign(self, path: str, expiration: int = 3600, **kwargs) -> str:
329
+ return cast(str, self._sync_fs.sign(path, expiration=expiration, **kwargs))
330
+
331
+ def checksum(self, path: str, **kwargs) -> int:
332
+ return cast(int, self._sync_fs.checksum(path, **kwargs))
333
+
334
+ def created(self, path: str) -> "datetime":
335
+ return self._sync_fs.created(path)
336
+
337
+ def modified(self, path: str) -> "datetime":
338
+ return self._sync_fs.modified(path)
339
+
340
+ def invalidate_cache(self, path: Optional[str] = None) -> None:
341
+ self._sync_fs.invalidate_cache(path)
342
+
343
+ async def _touch(self, path: str, truncate: bool = True, **kwargs) -> None:
344
+ await asyncio.to_thread(self._sync_fs.touch, path, truncate=truncate, **kwargs)
345
+
346
+
347
+ class AioS3File(S3File):
348
+ """Async-aware S3 file handle using ``S3AioExecutor``.
349
+
350
+ Functionally identical to ``S3File``; exists as a distinct type for
351
+ ``isinstance`` checks and to document the async execution model.
352
+ All parallel operations (range reads, multipart uploads) are dispatched
353
+ through the ``S3Executor`` interface — the ``S3AioExecutor``
354
+ provided by ``AioS3FileSystem`` uses the event loop instead of threads.
355
+ """
356
+
357
+ pass
@@ -0,0 +1,89 @@
1
+ # -*- coding: utf-8 -*-
2
+ from __future__ import annotations
3
+
4
+ import asyncio
5
+ from abc import ABCMeta, abstractmethod
6
+ from concurrent.futures import Future
7
+ from concurrent.futures.thread import ThreadPoolExecutor
8
+ from typing import Any, Callable, Optional, TypeVar
9
+
10
+ T = TypeVar("T")
11
+
12
+
13
+ class S3Executor(metaclass=ABCMeta):
14
+ """Abstract executor for parallel S3 operations.
15
+
16
+ Defines the interface used by ``S3File`` and ``S3FileSystem`` for submitting
17
+ work to run in parallel and for shutting down the executor when done.
18
+ Both ``submit`` and ``shutdown`` mirror the ``concurrent.futures.Executor``
19
+ interface so that ``as_completed()`` and ``Future.cancel()`` work unchanged.
20
+ """
21
+
22
+ @abstractmethod
23
+ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]:
24
+ """Submit a callable for execution and return a Future."""
25
+ ...
26
+
27
+ @abstractmethod
28
+ def shutdown(self, wait: bool = True) -> None:
29
+ """Shut down the executor, freeing any resources."""
30
+ ...
31
+
32
+ def __enter__(self) -> "S3Executor":
33
+ return self
34
+
35
+ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
36
+ self.shutdown(wait=True)
37
+
38
+
39
+ class S3ThreadPoolExecutor(S3Executor):
40
+ """Executor that delegates to a ``ThreadPoolExecutor``.
41
+
42
+ This is the default executor used by ``S3File`` and ``S3FileSystem``
43
+ for synchronous parallel operations.
44
+ """
45
+
46
+ def __init__(self, max_workers: int) -> None:
47
+ self._executor = ThreadPoolExecutor(max_workers=max_workers)
48
+
49
+ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]:
50
+ return self._executor.submit(fn, *args, **kwargs)
51
+
52
+ def shutdown(self, wait: bool = True) -> None:
53
+ self._executor.shutdown(wait=wait)
54
+
55
+
56
+ class S3AioExecutor(S3Executor):
57
+ """Executor that schedules work on an asyncio event loop.
58
+
59
+ Uses ``asyncio.run_coroutine_threadsafe(asyncio.to_thread(fn), loop)`` to
60
+ dispatch blocking functions onto the event loop's thread pool, returning
61
+ ``concurrent.futures.Future`` objects that are compatible with
62
+ ``as_completed()`` and ``Future.cancel()``.
63
+
64
+ This avoids thread-in-thread nesting when ``S3File`` is used from within
65
+ ``asyncio.to_thread()`` calls (the pattern used by ``AioS3FileSystem``).
66
+
67
+ Args:
68
+ loop: A running asyncio event loop.
69
+
70
+ Raises:
71
+ RuntimeError: If the event loop is not running when ``submit`` is called.
72
+ """
73
+
74
+ def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None) -> None:
75
+ self._loop = loop
76
+
77
+ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]:
78
+ if self._loop is not None and self._loop.is_running():
79
+ return asyncio.run_coroutine_threadsafe(
80
+ asyncio.to_thread(fn, *args, **kwargs), self._loop
81
+ )
82
+ raise RuntimeError(
83
+ "S3AioExecutor requires a running event loop. "
84
+ "Use S3ThreadPoolExecutor for synchronous usage."
85
+ )
86
+
87
+ def shutdown(self, wait: bool = True) -> None:
88
+ # No resources to release — work is dispatched to the event loop.
89
+ pass
@@ -5,6 +5,8 @@ import logging
5
5
  from io import TextIOWrapper
6
6
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
7
7
 
8
+ from fsspec import AbstractFileSystem
9
+
8
10
  from pyathena.converter import Converter
9
11
  from pyathena.error import OperationalError, ProgrammingError
10
12
  from pyathena.filesystem.s3 import S3FileSystem
@@ -62,6 +64,7 @@ class AthenaS3FSResultSet(AthenaResultSet):
62
64
  retry_config: RetryConfig,
63
65
  block_size: Optional[int] = None,
64
66
  csv_reader: Optional[CSVReaderType] = None,
67
+ filesystem_class: Optional[Type[AbstractFileSystem]] = None,
65
68
  **kwargs,
66
69
  ) -> None:
67
70
  super().__init__(
@@ -77,6 +80,7 @@ class AthenaS3FSResultSet(AthenaResultSet):
77
80
  self._arraysize = arraysize
78
81
  self._block_size = block_size if block_size else self.DEFAULT_BLOCK_SIZE
79
82
  self._csv_reader_class: CSVReaderType = csv_reader or AthenaCSVReader
83
+ self._filesystem_class: Type[AbstractFileSystem] = filesystem_class or S3FileSystem
80
84
  self._fs = self._create_s3_file_system()
81
85
  self._csv_reader: Optional[Any] = None
82
86
 
@@ -92,9 +96,9 @@ class AthenaS3FSResultSet(AthenaResultSet):
92
96
  if not self._csv_reader and not self._rows and pre_fetched_rows:
93
97
  self._rows.extend(pre_fetched_rows)
94
98
 
95
- def _create_s3_file_system(self) -> S3FileSystem:
99
+ def _create_s3_file_system(self) -> AbstractFileSystem:
96
100
  """Create S3FileSystem using connection settings."""
97
- return S3FileSystem(
101
+ return self._filesystem_class(
98
102
  connection=self.connection,
99
103
  default_block_size=self._block_size,
100
104
  )
@@ -178,7 +178,7 @@ exclude = [
178
178
  legacy_tox_ini = """
179
179
  [tox]
180
180
  isolated_build = true
181
- envlist = py{310,311,312,313,314}
181
+ envlist = py{310,311,312,313,314}-{pyathena,sqla,sqla_async}
182
182
 
183
183
  [gh-actions]
184
184
  python =
@@ -188,6 +188,12 @@ python =
188
188
  3.13: py313
189
189
  3.14: py314
190
190
 
191
+ [gh-actions:env]
192
+ TEST_TYPE =
193
+ pyathena: pyathena
194
+ sqla: sqla
195
+ sqla_async: sqla_async
196
+
191
197
  [testenv]
192
198
  allowlist_externals =
193
199
  uv
@@ -195,9 +201,9 @@ allowlist_externals =
195
201
  make
196
202
  commands =
197
203
  uv sync --group dev
198
- make test
199
- make test-sqla
200
- make test-sqla-async
204
+ pyathena: make test
205
+ sqla: make test-sqla
206
+ sqla_async: make test-sqla-async
201
207
  passenv =
202
208
  TOXENV
203
209
  AWS_*
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes