【AWS Athena】クエリ実行結果のデータを高速にPandas.DataFrameとして取得する

# はじめに
[PyAthena](https://pypi.org/project/pyathena/)を使うとPython3上でSQLを書き、それを実行してAthenaでのクエリ実行結果のデータを取得、Pandas.DataFrameで展開することができます。
ですが、取得する対象のデータが大きくなると、とても時間が掛かるので良い方法が無いか探してみました。


# 環境・使用ライブラリ
- windows 10 Home Edition
- Python 3.7.6
  - Jupyter Notebook (jupyter 1.0.0)
  - pip 20.3.3
  - s3fs 0.4.2
  - pandas 1.2.0
  - pyathena 1.10.0
  - boto3 1.16.63


# 結果
クエリ実行結果が50000000行 x 10列のデータを**1分6秒**で取得することができました。
PyAthena connect関数を使用しpd.read_sql()した場合は**4時間11秒**でした。

## 20210227 更新
[Dask.DataFrame](https://docs.dask.org/en/latest/dataframe.html)として取得した場合は**57.6秒**でした。PCメモリ容量が足りない場合はこちらのほうが良さそうです。
Daskである程度操作したあとに、compute()するのが良いと思います。
コードは[Gistで公開](https://gist.github.com/nijigen-plot/a325d310529752c4c272fb2dc9ff1bb2)しています。

## 20211127 更新

[awswrangler](https://pypi.org/project/awswrangler/)を使えばこんなことしなくても同等に早くなりました(**50.8秒**)
[【AWS Athena】クエリ実行結果を高速に取得する v2【Python】](https://qiita.com/nijigen_plot/items/07e5ee737804111e695c)

# PyAthenaについて
[PyAthena](https://pypi.org/project/pyathena/) はPython DB API 2.0 に準拠したAmazon Athena用のクライアントです。
PyAthenaを使うと、Athenaコンソール上でできる操作がPython上でできるようになります。
クエリを書いて実行したデータを取得したり、TABLEやVIEWを作ったりできます。

```python
import pandas as pd
from pyathena import connect

# AWS アクセスキー、シークレットキーを使用してAthenaへ接続
conn = connect(
    aws_access_key_id = 'アクセスキー',
    aws_secret_access_key = 'シークレットキー',
    s3_staging_dir = 'クエリ実行結果保管URL',
    region_name = 'リージョン'
    )
```

```python
# conn変数を使用してクエリの結果をpd.read_sqlでDataFrame化
%%time
query = """
SELECT * FROM "main_data"."riiid_train_data"
limit 10
"""

df = pd.read_sql(query, conn)
display(df)
```

実行結果
![SnapCrab_NoName_2021-2-14_13-28-14_No-00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/635079/2a2d87e7-2d5d-882b-69d7-c4f306768893.png)

上記画像の様にクエリの実行結果をDataFrameとして取得することができました。
このように、Athenaコンソール画面を使わなくても、データを取得してここから分析や集計等することができます。とても便利です。

## 取得するデータ量が多い場合
limit句の数字を`50000000`にして、再度取得してみます。
10行取得→50000000行取得になるので、とても大きいデータになります。

```python
%%time
query = """
SELECT * FROM "main_data"."riiid_train_data"
limit 50000000 -- 10から50000000に変更
"""

df = pd.read_sql(query, conn)
display(df)
```

![SnapCrab_NoName_2021-2-14_16-30-15_No-00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/635079/dd84dd48-f05b-f783-dc29-300297113b89.png)


なんと**4時間11秒**かかりました。とても遅いです。
ちなみに、PyAthenaから実行したクエリも**コンソールの`履歴`タブから見れます**ので、履歴を追って`Download the results in CSV format`ボタンからデータを落としてみましたが、ダウンロード時間は**約8分**と表示されました。


![SnapCrab_Athena - Google Chrome_2021-2-14_14-10-22_No-00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/635079/cd0d92b8-d260-afd5-c2f5-98290dc7861a.png)

どの部分が問題で、所要時間の乖離が起きているのか詳しくは分からないですが、個人的にはクエリ実行結果が**数百万行あたりから取得時間の遅さが目立って**きます。


# 高速に取得する方法
S3からクエリ実行結果のデータを直接取得します。

## テーブルを作成する
実行するクエリのテーブルを作成します。
[取得するデータ量が多い場合](#取得するデータ量が多い場合)のクエリをCREATE TABLE ASして、まずはテーブルを作ります

```python
cursor = connect(
    aws_access_key_id = 'アクセスキー',
    aws_secret_access_key = 'シークレットキー',
    s3_staging_dir = 'クエリ実行結果保管URL',
    region_name = 'リージョン'
    ).cursor() # conn変数と違うのは、最後に.cursor()が付きます。

# テーブルを作成
cursor.execute("""
CREATE TABLE main_data.fifty_million_data -- fifly_million_dataという名前でテーブルを作成する
WITH (format ='PARQUET') -- デフォルトでPARQUETですがわかりやすいように
AS SELECT * FROM "main_data"."riiid_train_data" limit 50000000 -- クエリは先程と同じ
""")

# テーブルプロパティを確認
query = """
SHOW CREATE TABLE main_data.fifty_million_data
"""
p_df = pd.read_sql(query, conn)
display(p_df)
```

テーブルプロパティ
![SnapCrab_NoName_2021-2-14_14-30-4_No-00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/635079/3db780fb-7bd8-aa26-dbba-1467ff9266f5.png)
テーブルプロパティは、コンソールでも確認できます。
コンソールで指定したテーブルの3つ点マークをクリック→プロパティの表示
![SnapCrab_Athena - Google Chrome_2021-2-14_14-31-4_No-00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/635079/8e71dcc6-15f4-2d90-87d5-34ebd6cb20db.png)

`場所`のURLにCREATE TABLE AS のクエリ実行結果のデータが入っています。
![SnapCrab_Athena - Google Chrome_2021-2-14_14-31-18_No-00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/635079/21dfb374-7d5e-bf24-07f4-891af903d805.png)

### S3のデータを見てみる
`場所`に記載されたURLにアクセスすると、manifest.csv,metadata,フォルダが格納されています。
![SnapCrab_NoName_2021-2-14_13-32-47_No-00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/635079/6fd0c3b3-7372-10d9-a535-cf3bc1c37a55.png)

フォルダ内を見てみると、たくさんのparquetデータがあります。（拡張子の表示がないですがparquetです）
クエリの実行結果のデータは30個のparquetデータに分けられて、S3に保存されてるみたいです。
![SnapCrab_NoName_2021-2-14_13-33-6_No-00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/635079/8bd6de75-1a0a-f4fa-a1ec-3162ee79ad9c.png)


## S3に保存されたデータを取得する
S3に保存されたデータを取得するには、まずそのURLが必要です。
[先程も書きました](#テーブルを作成する)が `SHOW CREATE TABLE`でテーブルのプロパティデータを取得します。

```python
query = """
SHOW CREATE TABLE main_data.fifty_million_data
"""
p_df = pd.read_sql(query, conn)

# LOCATIONの次のindexのデータを取得して、空白やシングルクォーテーションを除きます
data_location = p_df.iat[(p_df[p_df['createtab_stmt'] == 'LOCATION'].index + 1)[0], 0]
data_location = data_location.replace(' ','').replace("'","")
print(data_location)
```
data_locationの中身は`'s3://query-results-bucker/primary_folder/tables/f8d91c0a-3450-48e3-8fdb-9cfb23df30bd/'`です。
作成したテーブルのデータが保存されている場所ですね。

urlparseを使用して、netlocとpathに分けます。

```python
import urllib
scheme, netloc, path, params, query, fragment = urllib.parse.urlparse(data_location)
print(netloc)
print(path)
```
`query-results-bucker`(netloc)
`/primary_folder/tables/490c0c37-5e0e-4bdf-9590-37b6a4f19204/`(path)

boto3でS3へアクセス、ioでデータをPandas.read_parquetで読めるようにして、データを取得します。

```python
import boto3
import io
def read_s3_multiple_parquets(filepath, bucket_name, athena_access_key, athena_secret_key):
    # boto3.client,boto3.resourceにそれぞれアクセスキー、シークレットキーを入れて変数にします。
    s3_client = boto3.client(service_name = 's3', aws_access_key_id = athena_access_key, aws_secret_access_key = athena_secret_key)
    s3_resource = boto3.resource(service_name = 's3', aws_access_key_id = athena_access_key, aws_secret_access_key = athena_secret_key)
    # 指定されたディレクトリからファイル名を取得します。30個のparquetファイルのことです。
    s3_keys = [item.key for item in s3_resource.Bucket(bucket_name).objects.filter(Prefix = filepath)]
    # それぞれpd.read_parquetして、DataFrame化します。
    dfs = []
    for key in s3_keys:
        obj = s3_client.get_object(Bucket = bucket_name, Key = key)
        dfs.append(pd.read_parquet(io.BytesIO(obj['Body'].read())))
    # すべて結合した状態で返します。
    return pd.concat(dfs)

# read_s3_multiple_parquets関数を実行します。
df = read_s3_multiple_parquets(path[1:], netloc, athena_key.access(), athena_key.secret())
display(df)
```

取得できました。かかった時間は**1分6秒**です。
![SnapCrab_NoName_2021-2-14_16-48-46_No-00.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/635079/7b72e949-e10f-1674-896e-b430b6da542c.png)




## データを削除する
最後に、CREATE TABLE AS して作成したテーブルと、テーブルのために作成されたS3のデータを削除します。


```python
# テーブルを削除
cursor.execute("""
DROP TABLE main_data.fifty_million_data
""")

# S3上のデータは消えないので、それも消します
import s3fs
fs = s3fs.S3FileSystem(key = 'アクセスキー', secret = 'シークレットキー')
# データフォルダと、metadata, manifest.csvを削除
fs.rm(data_location, recursive = True)
fs.rm(data_location[:-1] + '.metadata')
fs.rm(data_location[:-1] + '-manifest.csv')
```
## 所要時間結果
[テーブルを作成する](#テーブルを作成する)～[データを削除する](#データを削除する)までの実行時間で、**1分6秒**になりました。
すべてを連結して時間を計測したコードは[Gistに乗せています](https://gist.github.com/nijigen-plot/22297a02b826056d01ba3e67d6cac5b1)ので、気になる方は見てみてください。

# 最後に
大幅に短縮できたので個人的には嬉しいですが、もっとスマートなやり方があれば教えてほしいですm(__)m