Write Athena Query Results in Parquet, Avro, ORC, JSON formats#
Keywords: AWS, Athena, Parquet, Avro, ORC, JSON, S3
Introduction#
Athena ้ป่ฎคไฝฟ็จ CSV ๆไปถๆ ผๅผๆฅไฟๅญ query result. ็ฑไบๅจๅคงๆฐๆฎ้ขๅๆๅพๅคๆดไผ็ง็ๆฐๆฎๆ ผๅผๆด้ๅไฟๅญๆฐๆฎ. ไป 2021-08-05 ๅนด่ตท, AWS Athena ๅผๅงๆฏๆ Parquet, Avro, ORC, JSON ็ญๆ ผๅผ. ไฝๆฏไฝฟ็จ่ฟไบๆ ผๅผ็ๆนๆณๅฎๆนๆๆกฃ่ฏด็ๅนถไธๆธ ๆฅ. ๆฌๆๅฐไป็ปๅฆไฝไฝฟ็จ Parquet ๆ ผๅผๆฅไฟๅญ Athena ็ๆฅ่ฏข็ปๆ (ๅ ถไปๆ ผๅผ้ฝ็ฑปไผผ).
UNLOAD command in Athena ๆฏไธไธช SQL ๅฝไปค, ๅฏไปฅ็จๆฅๆๅฎๅฐ query result ๅฏผๅบๅฐ S3. ่ไธๅฎๆฏๆไธไบๅๆฐ, ็จๆฅๆๅฎๅฏผๅบๆ ผๅผไปฅๅ่ฟไบๆ ผๅผ็่ฏฆ็ป้ ็ฝฎ.
ไฝ ๅจ UNLOAD ็ๆถๅ้่ฆๆๅฎไธไธช่ฟไธๅญๅจ็ S3 folder, ็จไบไฟๅญ parquet ๆไปถ. ็ฑไบ Athena ๆฏไธไธชๅนถ่กๅผๆ, ๆไปฅไผๅฏผๅบๅคไธชๆฐๆฎๆไปถ. ไฝ ๅจ่ฐ็จ boto3.athena_client.start_query_execution() API ็ๆถๅ้่ฆๆๅฎไธไธช S3 folder ็จๆฅไฟๅญ query result metadata. ่ฟไธชๆไปถๅคนๅไนๅ้ฃไธชๅฏไปฅ็ธๅไนๅฏไปฅไธๅ. ไฝๆๆจ่็จไธๅ็, ไปฅๆนไพฟๅบๅๆฐๆฎๅ metadata. ่ฟ้้ข็ metadata ๆไปถไธญๅ
ๆฌไธไธช ${query_execution_id}-manifest.csv ๆไปถ, ่ฎฐๅฝไบๅฏผๅบ็ๆๆ parquet ๆไปถ็ S3 URI ๅ่กจ.
ๆไปฅๆป็ปไธๆฅ, ไฝ ้่ฆๅ่ฟไนๅ ไปถไบ:
ๆๅๆฌ็ SQL ่ฏญๅฅ็จ UNLOAD ๅฐ่ฃ . ไพๅฆๅฆๆๅๆฅ็่ฏญๅฅๆฏ
SELECT * FROM table LIMIT 10, ้ฃไนๅฐ่ฃ ๅ็่ฏญๅฅๅฐฑๆฏUNLOAD (SELECT * FROM table LIMIT 10) TO '{result_s3_folder_uri}' WITH ( format = 'parquet' ).่ฐ็จ
start_query_execution()API ๆฅๆง่กๆฅ่ฏข.็จ Job Poll ๆจกๅผๆฏ้ๅ ็งๅฐฑๅปๆฅไธไธ execution status, ๅฆๆๆๅไบๅฐฑ่ฟ่กไธไธๆญฅ.
ไป metadata ๆไปถไธญ่ฏปๅ parquet ๆไปถ็ S3 URI ๅ่กจ.
ไป parquet ๆไปถไธญ่ฏปๅ dataframe ็ถๅๆผๆฅๆไธไธช.
ไธ้ขๆไธไธช่ๆฌๅฐไธ้ข็้ป่พๅฐ่ฃ ๆไบ็จๆทๅๅฅฝ็ๅฝๆฐ, ๅฏไพๅ่.
1# -*- coding: utf-8 -*-
2
3"""
4This module allow user to use parquet file to store Athena query result.
5It returns the athena query result as a
6`polars.LazyFrame <https://pola-rs.github.io/polars/py-polars/html/reference/lazyframe/index.html>`_.
7
8Requirements:
9
10 polars
11 s3fs
12 pyarrow
13
14You can run ``pip install polars s3fs pyarrow`` to install them.
15
16Usage::
17
18 from aws_athena_query import (
19 run_athena_query,
20 read_athena_query_result,
21 wait_athena_query_to_succeed,
22 )
23"""
24
25import typing as T
26import uuid
27import time
28import textwrap
29
30import s3fs
31import polars as pl
32import pyarrow.dataset
33
34
35if T.TYPE_CHECKING:
36 from boto_session_manager import BotoSesManager
37 from s3pathlib import S3Path
38
39
40def wait_athena_query_to_succeed(
41 bsm: "BotoSesManager",
42 exec_id: str,
43 delta: int = 1,
44 timeout: int = 30,
45):
46 """
47 Wait a given athena query to reach ``SUCCEEDED`` status. If failed, raise
48 ``RuntimeError`` immediately. If timeout, raise ``TimeoutError``.
49
50 .. versionadded:: 0.11.1
51 """
52 # ref: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/athena/client/get_query_execution.html
53 elapsed = 0
54 for _ in range(999999):
55 res = bsm.athena_client.get_query_execution(
56 QueryExecutionId=exec_id,
57 )
58 status = res["QueryExecution"]["Status"]["State"]
59 if status == "SUCCEEDED":
60 return
61 elif status in ["FAILED", "CANCELLED"]:
62 raise RuntimeError(f"execution {exec_id} reached status: {status}")
63 else:
64 time.sleep(delta)
65 elapsed += delta
66 if elapsed > timeout:
67 raise TimeoutError(f"athena query timeout in {timeout} seconds!")
68
69
70def _get_dataset_and_metadata_s3path(
71 s3dir_result: "S3Path",
72) -> T.Tuple["S3Path", "S3Path"]:
73 # the dataset folder uri will be used in the UNLOAD command
74 # the final parquet files will be stored in this folder
75 # note that this folder has to be NOT EXISTING before the execution
76 s3dir_dataset = s3dir_result.joinpath("dataset", uuid.uuid4().hex).to_dir()
77
78 # the metadata folder will be used to store the query result metadata file
79 # the metadata file will tell you the list of data file uris
80 s3dir_metadata = s3dir_result.joinpath("metadata").to_dir()
81 return s3dir_dataset, s3dir_metadata
82
83
84def read_athena_query_result(
85 bsm: "BotoSesManager",
86 s3dir_result: "S3Path",
87 exec_id: str,
88 verbose: bool = True,
89) -> pl.LazyFrame:
90 """
91 Load the athena query result from s3. The query has to be succeeded already.
92
93 :param bsm: boto_session_manager.BotoSesManager object.
94 :param s3dir_result: an S3Path object that represent a s3 directory to store
95 the athena query result.
96 :param exec_id: athena query execution id, you can get it from the
97 :func:`run_athena_query` function.
98 :param verbose: do you want to print the log?
99
100 :return: the lazy DataFrame of the result, If you just need to return the
101 regular DataFrame, you can do ``df = read_athena_query_result(...).collect()``.
102
103 .. versionadded:: 0.11.1
104 """
105 s3dir_dataset, s3dir_metadata = _get_dataset_and_metadata_s3path(s3dir_result)
106 # read the manifest file to get list of parquet file uris
107 # ref: https://pola-rs.github.io/polars/py-polars/html/reference/api/polars.scan_pyarrow_dataset.html
108 s3path_manifest = s3dir_metadata.joinpath(f"{exec_id}-manifest.csv")
109 s3uri_list = s3path_manifest.read_text(bsm=bsm).splitlines()
110
111 if verbose:
112 query_editor_console_url = (
113 f"https://{bsm.aws_region}.console.aws.amazon.com/athena"
114 f"/home?region={bsm.aws_region}#/query-editor/history/{exec_id}"
115 )
116 print(f"preview query in athena editor: {query_editor_console_url}")
117 print(f"query result manifest: {s3path_manifest.console_url}")
118 print(f"query result data: {s3dir_dataset.console_url}")
119 print(f"number of files in result: {len(s3uri_list)}")
120
121 if isinstance(bsm.profile_name, str):
122 file_system = s3fs.S3FileSystem(profile=bsm.profile_name)
123 else:
124 credential = bsm.boto_ses.get_credentials()
125 file_system = s3fs.S3FileSystem(
126 key=credential.access_key,
127 secret=credential.secret_key,
128 token=credential.token,
129 )
130 dataset = pyarrow.dataset.dataset(s3uri_list, filesystem=file_system)
131 lazy_df = pl.scan_pyarrow_dataset(dataset)
132 df = lazy_df.select(pl.col("*"))
133 return df
134
135
136def run_athena_query(
137 bsm: "BotoSesManager",
138 s3dir_result: "S3Path",
139 sql: str,
140 database: str,
141 catalog: T.Optional[str] = None,
142 compression: str = "gzip",
143 encryption_configuration: T.Optional[dict] = None,
144 expected_bucket_owner: T.Optional[str] = None,
145 acl_configuration: T.Optional[dict] = None,
146 workgroup: T.Optional[str] = None,
147 execution_parameters: T.Optional[T.List[T.Any]] = None,
148 result_cache_expire: T.Optional[int] = None,
149 client_request_token: T.Optional[str] = None,
150 delta: int = 1,
151 timeout: int = 10,
152 verbose: bool = True,
153) -> T.Tuple[pl.LazyFrame, str]:
154 """
155 Run athena query and get the result as a polars.LazyFrame.
156 With LazyFrame, you can do further select, filter actions before actually
157 reading the data, and leverage the parquet predicate pushdown feature to
158 reduce the amount of data to be read. If you just need to return the
159 regular DataFrame, you can do ``df = run_athena_query(...)[0].collect()``.
160
161 Example::
162
163 >>> from boto_session_manager import BotoSesManager
164 >>> from s3pathlib import S3Path
165 >>>
166 >>> bsm = BotoSesManager(profile_name="your_aws_profile")
167 >>> bucket = f"{bsm.aws_account_id}-{bsm.aws_region}-data"
168 >>> prefix = f"athena/results/"
169 >>> s3dir_result = S3Path(f"s3://{bucket}/{prefix}").to_dir()
170
171 >>> database = "your_database"
172 >>> sql = f"SELECT * FROM {database}.your_table LIMIT 10;"
173 >>> lazy_df, exec_id = run_athena_query(
174 ... bsm=bsm,
175 ... s3dir_result=s3dir_result,
176 ... sql=sql,
177 ... database=database,
178 ... )
179 >>> df = lazy_df.collect()
180 >>> df
181 ...
182
183 start_query_execution doc: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/athena/client/start_query_execution.html
184
185 :param bsm: boto_session_manager.BotoSesManager object.
186 :param s3dir_result: an S3Path object that represent a s3 directory to store
187 the athena query result.
188 :param sql: sql query string.
189 :param database: database name.
190 :param catalog: see start_query_execution doc
191 :param compression: compression algorithm to use when writing parquet file,
192 see all options here: https://docs.aws.amazon.com/athena/latest/ug/compression-formats.html
193 :param compression: see start_query_execution doc
194 :param encryption_configuration: see start_query_execution doc
195 :param expected_bucket_owner: see start_query_execution doc
196 :param acl_configuration: see start_query_execution doc
197 :param workgroup: see start_query_execution doc
198 :param execution_parameters: see start_query_execution doc
199 :param result_cache_expire: cache query result for X minutes, if None,
200 it means no cache, see more information here: https://docs.aws.amazon.com/athena/latest/ug/reusing-query-results.html
201 this is very helpful to reduce the cost.
202 :param client_request_token: see start_query_execution doc
203 :param delta: sleep time in seconds between each query status check.
204 :param timeout: timeout in seconds.
205 :param verbose: do you want to print the log?
206
207 :return: the tuple of two item, the first item is the lazy DataFrame of the result,
208 the second item is the athena query execution id (str)>
209
210 .. versionadded:: 0.11.1
211 """
212 # the sql query should not end with ;, it will be embedded in the final query
213 sql = sql.strip()
214 if sql.endswith(";"):
215 sql = sql[:-1]
216
217 s3dir_dataset, s3dir_metadata = _get_dataset_and_metadata_s3path(s3dir_result)
218
219 # ref: https://docs.aws.amazon.com/athena/latest/ug/unload.html
220 # use UNLOAD command to write result into data format other than csv
221 final_sql = textwrap.dedent(
222 f"""
223 UNLOAD ({sql})
224 TO '{s3dir_dataset.uri}'
225 WITH ( format = 'parquet', compression= '{compression}')
226 """
227 )
228
229 # ref: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/athena/client/start_query_execution.html
230 # you may need to customize this for more advanced query,
231 # like cross database query, federated query
232 query_execution_context = dict(Database=database)
233 if catalog:
234 query_execution_context["Catalog"] = catalog
235 result_configuration = dict(OutputLocation=s3dir_metadata.uri)
236 if encryption_configuration:
237 result_configuration["EncryptionConfiguration"] = encryption_configuration
238 if expected_bucket_owner:
239 result_configuration["ExpectedBucketOwner"] = expected_bucket_owner
240 if acl_configuration:
241 result_configuration["AclConfiguration"] = acl_configuration
242 kwargs = dict(
243 QueryString=final_sql,
244 QueryExecutionContext=query_execution_context,
245 ResultConfiguration=dict(
246 OutputLocation=s3dir_metadata.uri,
247 ),
248 )
249 if workgroup:
250 kwargs["WorkGroup"] = workgroup
251 if execution_parameters:
252 kwargs["Parameters"] = execution_parameters
253 if result_cache_expire:
254 kwargs["ResultReuseConfiguration"] = dict(
255 ResultReuseByAgeConfiguration=dict(
256 Enabled=True,
257 MaxAgeInMinutes=result_cache_expire,
258 )
259 )
260 if client_request_token:
261 kwargs["ClientRequestToken"] = client_request_token
262 res = bsm.athena_client.start_query_execution(**kwargs)
263
264 # the start_query_execution API is async, it returns the execution id
265 exec_id = res["QueryExecutionId"]
266
267 # wait for the execution to finish
268 wait_athena_query_to_succeed(bsm=bsm, exec_id=exec_id, delta=delta, timeout=timeout)
269 lazy_df = read_athena_query_result(
270 bsm=bsm,
271 s3dir_result=s3dir_result,
272 exec_id=exec_id,
273 verbose=verbose,
274 )
275 return lazy_df, exec_id
276
277
278if __name__ == "__main__":
279 import random
280 import textwrap
281
282 from boto_session_manager import BotoSesManager
283 from s3pathlib import S3Path
284
285 bsm = BotoSesManager()
286
287 bucket = f"{bsm.aws_account_id}-{bsm.aws_region}-data"
288 prefix = f"athena/results/"
289 s3dir_result = S3Path(f"s3://{bucket}/{prefix}").to_dir()
290
291 database = "dynamodb_to_datalake"
292 table = "transactions"
293
294 n = random.randint(100, 500)
295 sql = textwrap.dedent(
296 f"""
297 SELECT *
298 FROM "{database}"."{table}"
299 LIMIT {n};
300 """
301 )
302
303 lazy_df, exec_id = run_athena_query(
304 bsm=bsm,
305 s3dir_result=s3dir_result,
306 sql=sql,
307 database=database,
308 )
309 df = lazy_df.collect()
310 print(df.shape)
311 print(df)
312 assert df.shape[0] == n
Output:
query result manifest: https://console.aws.amazon.com/s3/object/111122223333-us-east-1-data?prefix=athena/results/metadata/9d59666e-0b1a-47d0-92ec-9c1627191ec4-manifest.csv
query result data: https://console.aws.amazon.com/s3/buckets/111122223333-us-east-1-data?prefix=athena/results/dataset/cfb91cdef7ea4c12877db974127d2966/
number of files in result: 10
(246, 18)
shape: (246, 18)
โโโโโโโโโโโโโโฌโโโโโโโโโโโโโฌโโโโโโโโโโโโโฌโโโโโโโโโโโโโฌโโโโโโฌโโโโโโโโโโโโโฌโโโโโโโโโโโฌโโโโโโโโโโโโฌโโโโโโโโโโโโโ
โ _hoodie_co โ _hoodie_co โ _hoodie_re โ _hoodie_pa โ ... โ create_mon โ create_d โ create_ho โ create_min โ
โ mmit_time โ mmit_seqno โ cord_key โ rtition_pa โ โ th โ ay โ ur โ ute โ
โ --- โ --- โ --- โ th โ โ --- โ --- โ --- โ --- โ
โ str โ str โ str โ --- โ โ str โ str โ str โ str โ
โ โ โ โ str โ โ โ โ โ โ
โโโโโโโโโโโโโโชโโโโโโโโโโโโโชโโโโโโโโโโโโโชโโโโโโโโโโโโโชโโโโโโชโโโโโโโโโโโโโชโโโโโโโโโโโชโโโโโโโโโโโโชโโโโโโโโโโโโโก
โ 2023080801 โ 2023080801 โ id:account โ create_yea โ ... โ 08 โ 07 โ 22 โ 32 โ
โ 5254041 โ 5254041_5_ โ :769-359-8 โ r=2023/cre โ โ โ โ โ โ
โ โ 0 โ 960,create โ ate_month= โ โ โ โ โ โ
โ โ โ _a... โ 08... โ โ โ โ โ โ
โโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโผโโโโโโโโโโโโผโโโโโโโโโโโโโค
โ 2023080801 โ 2023080801 โ id:account โ create_yea โ ... โ 08 โ 07 โ 22 โ 43 โ
โ 5254041 โ 5254041_9_ โ :333-843-8 โ r=2023/cre โ โ โ โ โ โ
โ โ 1 โ 563,create โ ate_month= โ โ โ โ โ โ
โ โ โ _a... โ 08... โ โ โ โ โ โ
โโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโผโโโโโโโโโโโโผโโโโโโโโโโโโโค
โ 2023080801 โ 2023080801 โ id:account โ create_yea โ ... โ 08 โ 07 โ 22 โ 43 โ
โ 5254041 โ 5254041_9_ โ :038-353-2 โ r=2023/cre โ โ โ โ โ โ
โ โ 2 โ 716,create โ ate_month= โ โ โ โ โ โ
โ โ โ _a... โ 08... โ โ โ โ โ โ
โโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโผโโโโโโโโโโโโผโโโโโโโโโโโโโค
โ 2023080801 โ 2023080801 โ id:account โ create_yea โ ... โ 08 โ 07 โ 22 โ 43 โ
โ 5254041 โ 5254041_9_ โ :554-923-0 โ r=2023/cre โ โ โ โ โ โ
โ โ 0 โ 842,create โ ate_month= โ โ โ โ โ โ
โ โ โ _a... โ 08... โ โ โ โ โ โ
โโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโผโโโโโโโโโโโโผโโโโโโโโโโโโโค
โ ... โ ... โ ... โ ... โ ... โ ... โ ... โ ... โ ... โ
โโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโผโโโโโโโโโโโโผโโโโโโโโโโโโโค
โ 2023080801 โ 2023080801 โ id:account โ create_yea โ ... โ 08 โ 07 โ 22 โ 43 โ
โ 5254041 โ 5254041_9_ โ :526-507-4 โ r=2023/cre โ โ โ โ โ โ
โ โ 59 โ 992,create โ ate_month= โ โ โ โ โ โ
โ โ โ _a... โ 08... โ โ โ โ โ โ
โโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโผโโโโโโโโโโโโผโโโโโโโโโโโโโค
โ 2023080801 โ 2023080801 โ id:account โ create_yea โ ... โ 08 โ 07 โ 22 โ 43 โ
โ 5254041 โ 5254041_9_ โ :555-852-7 โ r=2023/cre โ โ โ โ โ โ
โ โ 60 โ 716,create โ ate_month= โ โ โ โ โ โ
โ โ โ _a... โ 08... โ โ โ โ โ โ
โโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโผโโโโโโโโโโโโผโโโโโโโโโโโโโค
โ 2023080801 โ 2023080801 โ id:account โ create_yea โ ... โ 08 โ 07 โ 22 โ 43 โ
โ 5254041 โ 5254041_9_ โ :531-448-3 โ r=2023/cre โ โ โ โ โ โ
โ โ 61 โ 070,create โ ate_month= โ โ โ โ โ โ
โ โ โ _a... โ 08... โ โ โ โ โ โ
โโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโโโผโโโโโโผโโโโโโโโโโโโโผโโโโโโโโโโโผโโโโโโโโโโโโผโโโโโโโโโโโโโค
โ 2023080801 โ 2023080801 โ id:account โ create_yea โ ... โ 08 โ 07 โ 22 โ 43 โ
โ 5254041 โ 5254041_9_ โ :491-600-3 โ r=2023/cre โ โ โ โ โ โ
โ โ 62 โ 033,create โ ate_month= โ โ โ โ โ โ
โ โ โ _a... โ 08... โ โ โ โ โ โ
โโโโโโโโโโโโโโดโโโโโโโโโโโโโดโโโโโโโโโโโโโดโโโโโโโโโโโโโดโโโโโโดโโโโโโโโโโโโโดโโโโโโโโโโโดโโโโโโโโโโโโดโโโโโโโโโโโโโ