【Python】AthenaへSQLを発行しSELECT結果をデータフレームに取得する方法

AWSのAthenaでは、S3内のデータをPrestoベースのSQLを使用して参照できます。

フラットファイルにクエリで検索をかけられます。

私も業務でAthenaを使用していますが、PythonからAthenaへSQLを発行したいケースは多々あるため、方法を記載しておきます。以前はboto3などを利用していた実装していましたが、PyAthenaというライブラリを使用する方が簡単でした。

必要なライブラリをインストールします。

pip install pandas
pip install pyathena

DataFrameとして結果を取得するため、pandasもインストールします。

SQLの発行を実装します。

from pyathena import connect
import pandas as pd

# AWS接続情報
AWS_ACCESS_KEY = 'access_key'
AWS_SECRET_ACCESS_KEY = 'secret_access_key'
AWS_REGION_NAME = 'region_name'
AWS_S3_STAGING_DIR = 'S3://bucket-name/key-name'

cn = connect(aws_access_key_id=AWS_ACCESS_KEY,
             aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
             s3_staging_dir=AWS_S3_STAGING_DIR,
             region_name=AWS_REGION_NAME)

# 実行SQL
sql = "SELECT * FROM table"

# DataFrameとして取得
df = pd.read_sql(sql, cn)

AWS接続情報と実行SQLは、接続先環境にあわせて変更してください。上記のコードを実行するとSQLの結果をDataFrameとして取得できます。

続けて、DataFrameのメソッドを使用して加工、分析などできます。

DataFrameからそのままCSVファイルとして出力するためには、以下の記述を追加します。

df.to_csv('table.csv', encoding='utf-8')

PyAthenaを使用して、簡単にPythonプログラムからAthenaにSQLを発行できます。PyAthenaの実装が簡単すぎて、これを知ったときは衝撃でした。