AWS Lambda から Amazon Athena を呼び出すための手軽な小技を紹介します。AWS SDK for pandas の SQL 実行クライアントとしての邪道な使い方です。
なぜ Lambda から Athena を呼び出すのか
私は ELT パイプラインを運用しており、オーケストレーターとして AWS Step Functions、変換エンジンとして Athena を使っています。Athena は Step Functions と直接統合できますが、時間範囲の計算などの汎用プログラミング言語のロジックを実行する環境も別途必要です。そのために Lambda を使っています。
オーケストレーターと変換エンジンの他の選択肢も検討しました。例えば、オーケストレーターの Step Functions の代わりに dbt を使えば、Lambda から Athena を呼び出す必要はなくなります。しかし、私のテーブルの依存関係はシンプルすぎて dbt を導入するほどではありません。
AWS SDK for pandas とは
AWS が提供する Python ライブラリで、pandas と AWS のデータソースを統合するものです。本来の用途は pandas の DataFrame API を使ったデータ分析・加工です。
パッケージ名は awswrangler です。「AWS SDK for pandas」という名前は旧称からの改名で、PyPI 上のパッケージ名は旧称のままです。
しかし、私は DataFrame を一切使っていません。awswrangler.athena.start_query_execution 関数だけを SQL 実行クライアントとして使っています。以下の代替案と比べると awswrangler が手軽です。
- boto3 や AWS SDK for JavaScript のような基本的な AWS SDK を使う場合、低レベルな Athena API に対してポーリングループを自前で書く必要があります。一方、
awswranglerは高レベル API を提供しており、コードがシンプルになります。 -
PyAthena のようなサードパーティ製の Athena クライアントを使う場合、自前で Lambda Layer をパッケージングして保守する必要があります。一方、
awswranglerは AWS 管理の Lambda Layer として提供されているため、Lambda Layer の自作が不要です。
後述のように、15分の Lambda の実行時間の上限以内でクエリが終わらないリスクがある場合、Step Functions 側でのポーリングループが必要になります。この場合、この小技は適用できません。
コード例
import awswrangler as wr
sql = """
INSERT INTO t1(col1, col2)
SELECT col1, col2 FROM ...
WHERE col3 BETWEEN :begin AND :end
-- Python DB API 2.0 の named parameter 構文
"""
response = wr.athena.start_query_execution(
sql=sql,
database=my_glue_db, # Glue Catalog の名前空間
s3_output="s3://my-athena-query-result",
paramstyle="named",
params={"begin": ..., "end": ...},
wait=True,
)
wait は必ず明示的に True にします。そうしないと、 Query Execution ID を即座に返してしまい、結局 GetQueryExecution API をポーリングすることになってしまいます。
wait を True にすると、start_query_execution はクエリが成功または失敗するまで待機して結果を返します。response は GetQueryExecution API が返す辞書です。
response["Status"]["State"] は文字列で、"SUCCEEDED" のときだけ実行成功を意味します。
if response["Status"]["State"] != "SUCCEEDED":
raise RuntimeError(f"{response=}")
awswrangler.athena.start_query_execution の公式ドキュメントはこちら。
start_query_execution で実行できる SQL 文
AWS SDK for pandas の Athena クライアントは、おそらく boto3 の Athena クライアントのラッパーだと思います。したがって Athena サービスが受け付ける文であれば問題なく実行できるはずだと思います。これは推測です。
パイプラインで、実際に動いた経験で確認している範囲では、INSERT・MERGE・DELETE といった DML、さらに Apache Iceberg テーブルに対する OPTIMIZE ... REWRITE DATA USING BIN_PACK によるコンパクションと VACUUM を問題なく実行しています。
ただし VACUUM には1回の実行で削除できるファイル数の上限があるため、リトライロジックが必要です。ポイントは response["Status"]["StateChangeReason"] に "ICEBERG_VACUUM_MORE_RUNS_NEEDED" という部分文字列が含まれているかを見ることです。リトライ自体は Step Functions レベルで行うのが自然でしょう。こちらの記事が参考になります。
Athena より先に Lambda がタイムアウトしないか?
Lambda には15分という実行時間の上限があります。これが Athena のクエリ完了に十分かどうかは、データ量やテーブル・ファイル形式に依存します。
参考までに、私のテーブルは Iceberg で、Athena での処理はかなり高速です。夜間バッチパイプラインを運用しており、約2000万行のファクトを取り込み、ディメンションを更新し、ファクトとディメンションを結合・分類して約400万行に事前集計しています。その後コンパクションと VACUUM を実行します。各 Lambda 関数は10〜30秒で完了しています。
私の場合、タイムアウトのリスクはありません。もし、タイムアウトのリスクがあれば、 Query Execution ID を即座に取得し、Step Functions 側で GetQueryExecution API をポーリングするべきです。もちろん、その場合、AWS SDK for pandas を使う意味はなくなるので、boto3 で十分です。
API の安定性について
awswrangler の Lambda Layer をバージョン固定しているので、awswrangler 側の API が勝手に変わりません。また、前述のとおり response は GetQueryExecution API の返り値そのままです。Athena 側の API が変わるとしたら、それは AWS SDK for pandas を使っていようが boto3 を使っていようが、どの Athena クライアントを使っていても同じ問題です。