AWS Data Wranglerのto_parquetとread_sql_queryの使い方

# はじめに
はじめまして！[株式会社オークン](https://o-kun.com/)のUDです！

今回はあまり参考事例がなく、悪戦苦闘したAWS Data Wranglerの使い方について、もっとも使用頻度の高かったawswrangler.s3.to_parquetとawswrangler.athena.read_sql_queryの使い方を中心にまとめてみたいと思います。

# 使用したもの

+ メタデータカタログ：Glue Data Catalog
+ データストア：S3
+ Athena
+ AWS Data Wrangler
+ AWS CDK
+ direnv

※なお、テーブル定義の変更はAWS Data Wranglerでは行わない前提で、テーブル定義とデプロイはCDKで行いました。

# 公式ドキュメント
+ AWS Data Wranglerの[API Reference](https://aws-data-wrangler.readthedocs.io/en/stable/api.html#amazon-athena)
 
# テーブル設計
テーブル設計の観点においても、to_parquetとread_sql_query両方の側面から設計を考える必要がありました。
パフォーマンスと料金面を考慮してテーブル設計を進める上で、大事なポイントについては以下のページにわかりやすくまとめられています。

+ [Amazon Athena 初心者向けハンズオン](https://www.slideshare.net/AmazonWebServicesJapan/amazon-athena-81216877)
+ [Top 10 Performance Tuning Tips for Amazon Athena](https://aws.amazon.com/jp/blogs/big-data/top-10-performance-tuning-tips-for-amazon-athena/)

上記のページからは見えないつまずきポイントも含めて、テーブル設計で考慮した方が良い主要なポイントをまとめると、次のような内容になります。

|項目| Athena | S3 |
|:------|:------|:--|
|主にかかる料金|・スキャンしたデータ1TBにつき5ドル<br>※別リージョンからデータを読み込むと、S3のデータ転送量がかかるため注意|・ストレージの料金<br>・PUTの料金|
|AWS Data Wranglerで主に使用する関数|[awswrangler.athena.read_sql_query](https://aws-data-wrangler.readthedocs.io/en/stable/stubs/awswrangler.athena.read_sql_query.html#awswrangler-athena-read-sql-query)|[awswrangler.s3.to_parquet 等](https://aws-data-wrangler.readthedocs.io/en/stable/stubs/awswrangler.s3.to_parquet.html#awswrangler.s3.to_parquet)|
|列指向ファイル形式の活用|列指向フォーマットの活用し、SELECTする列を指定する事で、指定された列のみをスキャンすることができ、スキャン量を減らせる|ー|
|ファイル圧縮の活用|snappy等の形式にファイルを圧縮する事で、スキャンするデータ量を減らせる|圧縮することでストレージ料金が抑えられる|
|パーティションの活用|WHERE文でよく使用するカラムをパーティションにする事で、スキャン量が減らせる|・並列処理でto_parquetを使用している場合は、modeをappendにするか、overwrite_partitionsでパーティション設計を並列処理に耐えうるようにする必要がある<br>・パーティションが多いとPUT回数が増える<br>※to_parquetもパーティション別に1オブジェクトずつPUTする|
|オブジェクトサイズの目安|ー|・ファイルサイズが小さすぎる場合、読み取り処理のオーバーヘッドに時間を費やす事になるため、128MB以上が目安とされている<br>・S3の料金体系として、128KB以下のオブジェクトは一律で128KBとしてカウントされるため、小さすぎるともったいない|


# AWS Data Wranglerの使い方に関するポイント
## credentialsの読み込み
AWS Data Wranglerのあらゆる関数の引数で指定することができるboto3_sessionですが、その名の通り、実態は[boto3のSession](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html?highlight=session#module-boto3.session)です。
AWS Data Wranglerの各関数の引数でboto3_sessionを指定しない場合、デフォルトでセットされているSession(DEFAULT_SESSION)を参照します。

DEFAULT_SESSIONは内部的に生成されますが、生成プロセスは様々です。
credentialsが読み込まれる優先順位については、[こちら](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html#configuring-credentials)にまとめられています。

例えば、各関数の引数でboto3_sessionを指定せずとも、環境変数からaws_access_key_id、aws_secret_access_key、aws_session_token、aws_default_regionを読み込み、DEFAULT_SESSIONを生成することが可能です。

```
# 環境変数(.envrcの場合)
export AWS_ACCESS_KEY_ID=""
export AWS_SECRET_ACCESS_KEY=""
export AWS_SESSION_TOKEN=""
export AWS_DEFAULT_REGION=""
```

```
# awswrangler.athena.read_sql_query使用箇所
import awswrangler as wr

def sample_def():
    query = 'SELECT * FROM sample_table'

    wr.athena.read_sql_query(
        database='sample_database',
        sql=query,
        ctas_approach=False,
        # boto3_session=, ← 不要になる
        max_cache_seconds=0,
        max_cache_query_inspections=50,
        max_remote_cache_entries=50,
        max_local_cache_entries=100,
    )
```

## Global Configurationsの活用
AWS Data Wranglerのチュートリアルに記載のある[Global Configurations](https://aws-data-wrangler.readthedocs.io/en/stable/tutorials/021%20-%20Global%20Configurations.html)を使用すると、上記read_sql_queryの各種設定を共通化することが可能です。

```
import awswrangler as wr

def init_config():
    # database名
    wr.config.database = "sample_database"
    # CTASを有効にするか
    wr.config.ctas_approach = False
    # 以下、read_sql_query実行時の各種キャッシュに関する設定
    wr.config.max_cache_seconds = 0
    wr.config.max_cache_query_inspections = 50
    wr.config.max_remote_cache_entries = 50
    wr.config.max_local_cache_entries = 100

def sample_def():
    query = 'SELECT * FROM sample_table'

    wr.athena.read_sql_query(
        sql=query,
    )

if __name__ == "__main__":
    init_config()
    sample_def()
```


## read_sql_queryのctas_approachについて
[awswrangler.athena.read_sql_query](https://aws-data-wrangler.readthedocs.io/en/stable/stubs/awswrangler.athena.read_sql_query.html#awswrangler-athena-read-sql-query)のctas_approachをTrueにすると[CTAS](https://docs.aws.amazon.com/ja_jp/athena/latest/ug/ctas.html)が有効になり、一時テーブルがGlue上に作成され、クエリ結果がS3に他のストレージ形式 (Parquet、ORC など) として保存されます。
このクエリ結果を参照することによりAthenaのパフォーマンスが向上するため、クエリの結果が中〜大規模になる場合は高速となります。
ただし、一時テーブルをGlue上に作成する必要があるため、Glueでのテーブルの作成・削除権限が必要になります。

## to_parquetのmodeについて
MYSQL等のRDBに慣れている方は、RDBをイメージしながら実装を進めると痛い目を見ます…。何度データが消えたことか…。
データの実態はあくまでS3のオブジェクトですので、appendはファイルのcreate、overwrite系はファイルのdelete&createが行われている事がイメージできていると理解しやすいと思います。

| mode | 仕様・注意点 |
|:------:|:------|
| append |・to_parquetの度にバケットまたはパーティションにオブジェクトが追加される|
| overwrite |・バケットまたは全てのパーティションのオブジェクトをdeleteした後に、新しいオブジェクトを作成する<br>・パーティションを含む場合は、Glueのパーティション情報も更新される<br>※全てのパーティションのオブジェクトが一度deleteされるため、誤って消滅してしまわないように、実装・開発上は細心の注意が必要|
| overwrite_partitions |・対象のパーティション配下のオブジェクトをdeleteした後に、新しいオブジェクトを作成する<br>・パーティションを含む場合は、Glueのパーティション情報も更新される|

## to_parquetと合わせて使用したい関数
modeがoverwriteの場合、テーブル定義の変更もto_parquet時に行えてしまいますが、テーブルの定義変更をto_parquet時に行わない場合、[wr.catalog.get_table_location](https://aws-data-wrangler.readthedocs.io/en/stable/stubs/awswrangler.catalog.get_table_location.html#awswrangler.catalog.get_table_location)でデータセットのパスを取得し、[wr.catalog.get_table_types](https://aws-data-wrangler.readthedocs.io/en/stable/stubs/awswrangler.catalog.get_table_types.html#awswrangler.catalog.get_table_types)で列名とデータ型がセットになったdictを取得すると処理が楽になります。

例えば、とあるビデオストリーミングサービスで扱っているテーブルが次のような構造をしていた場合
![video_streaming_schema.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2403888/23206f54-89d5-3e6e-ffee-c3e3e486e62d.png)
to_parquetを行う場合は次のようなコードになります。

```
import awswrangler as wr
from datetime import datetime
import pandas as pd

def init_config():
    # database名
    wr.config.database = "sample_database"

def sample_to_parquet():
    # サンプルのdataframeを作成
    watched_date = datetime(2021, 12, 20).date()
    dict_watched_history = dict(
        title=['スターウォーズ', 'キングスマン', 'インセプション'],
        genre=['SF', 'アクション', 'サスペンス'],
        user_name=['太郎', 'Mr.SASUKE', 'UD'],
        watched_date=[watched_date, watched_date, watched_date]
    )
    df = pd.DataFrame(data=dict_watched_history)

    # テーブル名
    table_name = 'sample_watched_video_history'

    # datasetとして指定しているS3のパスを取得
    s3_dataset_location_path = wr.catalog.get_table_location(table=table_name)

    # テーブル定義からdtypeを取得
    dtype = wr.catalog.get_table_types(table=table_name)

    # to_parquet処理
    wr.s3.to_parquet(
        df=df, 
        path=s3_dataset_location_path,
        dataset=True,  
        table=table_name, 
        mode='overwrite_partitions', 
        dtype=dtype, 
        schema_evolution=False, 
        partition_cols=['watched_date'] 
    )

if __name__ == "__main__":
    init_config()
    sample_to_parquet()

```

# おわりに
上記は、クローラを使用しないかなり限定的な利用用途にはなりますが、同じような使い方をされる方の参考になれば幸いです。
次回は、to_parquetしたオブジェクトのロールバックについて、悪戦苦闘した内容を書きたいと思います。