Jupyter NotebookでAmazon Athenaを使う

はじめに

Amazon Athenaは大規模なデータに対して、標準的なSQLを利用して素早くクエリを実行できます。このため、探索的なデータの分析のように、インタラクティブにデータを操作したい場合などに重宝します。

インタラクティブに操作を繰り返すことで、最終的に求めていたデータを取得するSQLを手に入れることはできますが、時には記録・共有の面からその調査結果を過程を含めてまとめておきたい時があります。

Jupyter Notebookを用いて調査結果をまとめれば、ドキュメント内にコードを含めることができるため、これをそのまま実行可能な調査資料として活用することができます。

本記事では、Jupyter Notebookを用いてAmazon Athenaを操作する方法について紹介します。

準備

Jupyter NotebookでAmazon Athenaを利用するためには、以下の拡張が必要となります。

• ipython-sql
• psycopg2
• PyAthenaJDBC

以下のコマンドでインストールします。

$ pip install jupyter
$ pip install PyAthenaJDBC
$ pip install psycopg2

ipython-sqlについては、PyPIに提供されている最新(2017/12/03時点)のバージョン(0.3.8)ではauto commitの無効化が出来ず、クエリを実行した場合にエラーとなってしまうため、最新版をGitHubより取得し、インストールします。

$ git clone git@github.com:catherinedevlin/ipython-sql.git
$ cd ipython-sql
$ python setup.py install

psycopg2は今回は利用しませんが、インストールしていないとipython-sqlを実行した際にエラーとなってしますため、ひとまず入れています。

クエリの実行

実際にjupyter notebookからクエリを実行してみます。

Terminalから適当なディレクトリを作成し、jupyter notebookを起動します。

$ mkdir jupyter-athena-test
$ cd jupyter-athena-test
$ jupyter notebook .

ブラウザからJupyter Notebookが起動するので、適当にnotebookを作成し、以下を実行します。

%load_ext sql
%config SqlMagic.autocommit = False
conn_string = 'awsathena+jdbc://XXXXXXXXXXXXXX:XXXXXXXXXXXXXXXXXXXXXXXXXX@athena.ap-northeast-1.amazonaws.com:443/default?s3_staging_dir=s3://athena-output-buckets-masuwo3/'

JDBCからAthenaに接続するための接続文字列は以下のようなフォーマットになっています。

awsathena+jdbc://{access_key}:{secret_key}@athena.{region_name}.amazonaws.com:443/{schema_name}?s3_staging_dir={s3_staging_dir}

最後にSQLを実行して結果を取得してみます。

%%sql $conn_string
SELECT * FROM "sampledb"."cloudfront_logs" limit 10;

無事、チュートリアルで作成したcloudfront_logテーブルの中身が取得できました。

おまけ

Athenaの接続文字列はそのままでは読み書きには少し煩雑なので、以下のような関数を用意しておくと便利です。

import boto3

def gen_athena_conn(schema, s3_staging_dir, profile=None):
    """
    AthenaにJDBCから接続する際の接続文字列を生成する
    profileが指定されなければデフォルトが設定される
    """
    session = boto3.Session(profile_name=profile)
    credential = session.get_credentials()

    access_key = credential.access_key
    secret_key = credential.secret_key
    region_name = session.region_name

    templ = "awsathena+jdbc://{access_key}:{secret_key}@athena.{region_name}.amazonaws.com:443/"\
            "{schema_name}?s3_staging_dir={s3_staging_dir}"

    conn = templ.format(
        access_key=access_key,
        secret_key=secret_key,
        region_name=region_name,
        schema_name=schema,
        s3_staging_dir=s3_staging_dir)

    return conn

利用例は以下の通りです。

conn_string = gen_athena_conn(
    schema='sampledb',
    s3_staging_dir='s3://athena-output-buckets-masuwo3/')

%%sql $conn_string
SELECT * FROM "sampledb"."cloudfront_logs" limit 10;

参考

• ipython-sql
• PyAthenaJDBC