Livebook から Amazon Athena で S3 上の CSV ファイルに SQL クエリを実行する

## はじめに

Livebook　から AWS のサービスを操作するシリーズです

今回は Amazon Athena を操作します

Athena は S3 上に保存したテキストファイル（CSVやTSV、JSON）を RDB のテーブルに見立て、 SQL を実行できるサービスです

https://aws.amazon.com/jp/athena/?whats-new-cards.sort-by=item.additionalFields.postDateTime&whats-new-cards.sort-order=desc

SQL クエリによって 1 TB スキャンしたら $5 という金額なので、かなり安く利用できます

実は私もつい最近このサービスを知ったのですが、いちいちデータベースを作りたくないときに便利そうです

今回は Livebook から S3 上に CSV をアップロードし、 Athena から SQL クエリを実行してみます

実装したノートブックはこちら

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/aws/athena.livemd

## 事前作業

AWS のアカンウトと、 S3 の権限を持った IAM ユーザーと、その認証情報（ACCESS_KEY_ID と SECRET_ACCESS_KEY）が必要です

https://docs.aws.amazon.com/ja_jp/IAM/latest/UserGuide/id_users_create.html#id_users_create_console

## 実行環境

Livebook 0.7.2 の Docker　イメージを元にしたコンテナで動かしました

https://hub.docker.com/layers/livebook/livebook/0.7.2/images/sha256-23b577ff5c956fa017bc148d41e219d3697b159813feeefbfc37a6853ccaf62e?context=explore

コンテナ定義はこちらを参照

https://github.com/RyoWakabayashi/elixir-learning

## セットアップ

[前回](https://qiita.com/RyoWakabayashi/items/a6fd4d428851eec03bbf)は [ex_aws](https://github.com/ex-aws/ex_aws) を使いましたが、 [ex_aws_athena](https://github.com/vectordotdev/ex_aws_athena) の最終コミットが4年前でアーカイブされていたため、今回は aws-elixir の方を使ってみます

https://github.com/aws-beam/aws-elixir

README には以下のように書いてあります

> Most of the code is generated using the aws-codegen library from the JSON descriptions of AWS services provided by Amazon. They can be found in lib/aws/generated.

実は aws-elixir のコードは aws-sdk-go の仕様から自動生成されたものです

https://github.com/aws/aws-sdk-go

ほぼ全てのの AWS サービスをカバーしている代わりに、ドキュメントがあまり親切ではありません

ほとんどの関数の第2引数が `input` で、これだけ見ても何を渡せばいいのか分かりません

https://hexdocs.pm/aws/AWS.Athena.html

input の詳細な内容は各 API の仕様を見ましょう

https://docs.aws.amazon.com/ja_jp/athena/latest/APIReference/API_StartQueryExecution.html

```elixir
Mix.install([
  {:aws, "~> 0.13"},
  {:uuid, "~> 1.1"},
  {:hackney, "~> 1.18"},
  {:explorer, "~> 0.3"},
  {:kino, "~> 0.7"}
])
```

エイリアス等の準備をします

```elixir
alias Explorer.DataFrame
alias Explorer.Series
require Explorer.DataFrame
```

## 認証

入力エリアを用意し、そこに IAM ユーザーの認証情報を入力します

ACCESS_KEY_ID　と SECRET_ACCESS_KEY は秘密情報なので、値が見えないように `Kino.Input.password` を使います

https://qiita.com/RyoWakabayashi/items/debb1d048613ce505988

```elixir
access_key_id_input = Kino.Input.password("ACCESS_KEY_ID")
```

```elixir
secret_access_key_input = Kino.Input.password("SECRET_ACCESS_KEY")
```

リージョンもここで入力しておきましょう

```elixir
region_input = Kino.Input.text("REGION")
```

![スクリーンショット 2022-12-03 0.01.54.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/9c877571-862a-22bd-528f-32d5cb6c28b5.png)

各認証情報を使って、 aws-elixir で AWS にリクエストを投げるためのクライアントを生成します

```elixir
client =
  AWS.Client.create(
    Kino.Input.read(access_key_id_input),
    Kino.Input.read(secret_access_key_input),
    Kino.Input.read(region_input)
  )
```

![スクリーンショット 2022-12-05 14.21.03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/3b4a9ca7-50ac-b5c0-9375-851513544378.png)

## データ作成

今回は Explorer からワインデータセットを使用します

```elixir
wine_df = Explorer.Datasets.wine()

wine_df
|> Kino.DataTable.new()
```

![スクリーンショット 2022-12-05 14.49.01.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/cd9c729a-1532-b526-f356-bf3a2fb27788.png)

CSV に保存します

```elixir
csv_filename = "wine.csv"

wine_df
|> DataFrame.to_csv(csv_filename)
```

## データアップロード

アップロード先の S3 バケットを確認します

```elixir
buckets_res =
  client
  |> AWS.S3.list_buckets()
  |> elem(1)
```

![スクリーンショット 2022-12-05 14.59.49.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/8307c60b-6f5f-83fc-44b2-6691bf4ad350.png)

これでは見にくいので、データテーブルで表示します

```elixir
buckets_res["ListAllMyBucketsResult"]["Buckets"]["Bucket"]
|> Kino.DataTable.new()
```

![スクリーンショット 2022-12-05 15.01.03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/21a13b4e-fc4f-09d6-e20e-bc1a8d7f6a01.png)

バケット名を入力します

```elixir
bucket_name_input = Kino.Input.text("BUCKET_ANME")
```

![スクリーンショット 2022-12-05 15.03.56.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/3f1effe8-2c8d-9734-df00-97faa7fff057.png)

CSV アップロード先のバケットとパスを指定します

```elixir
bucket_name = Kino.Input.read(bucket_name_input)
athena_prefix = "athena/"
```

CSV　ファイルを S3 にアップロードします

```elixir
file = File.read!(csv_filename)
md5 = :crypto.hash(:md5, file) |> Base.encode64()

client
|> AWS.S3.put_object(
  bucket_name,
  athena_prefix <> csv_filename,
  %{"Body" => file, "ContentMD5" => md5}
)
```

![スクリーンショット 2022-12-05 15.04.40.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/59eec6ca-2db1-3720-48bc-6b7be23581c0.png)

## データベース定義

まずデータベースを作成します

（作成せずに default データベースを使っても良いです）

Athena では SQL を実行する度に一意のトークンを求められます

https://docs.aws.amazon.com/athena/latest/APIReference/API_StartQueryExecution.html#athena-StartQueryExecution-request-ClientRequestToken

AWS CLI の場合は内部的に生成してくれているのですが、 API を直接呼ぶ場合は自分で生成しないといけません

トークンは 32 〜 128 文字ということなので、 UUID を使います

```elixir
create_db_query = "CREATE DATABASE athena_sample"
create_db_token = UUID.uuid1()
```

![スクリーンショット 2022-12-05 15.08.53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/892f9cc1-c224-bc3d-f74d-e55b48ea7985.png)

クエリを実行します

`ResultConfiguration` の `OutputLocation` に S3 バケットを指定します（指定しないとエラーになります）

クエリは非同期に実行されるため、即結果が返ってきません

実行結果の QueryExecutionId を使って実行状態を取得する必要があります

```elixir
exec_id =
  client
  |> AWS.Athena.start_query_execution(%{
    "QueryString" => create_db_query,
    "ClientRequestToken" => create_db_token,
    "ResultConfiguration" => %{
      "OutputLocation" => "s3://#{bucket_name}"
    }
  })
  |> elem(1)
  |> then(& &1["QueryExecutionId"])
```

![スクリーンショット 2022-12-05 15.11.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/c6bf2097-bb79-fe3d-5c87-e84e1a152953.png)

実行状況を確認します

```elixir
client
|> AWS.Athena.get_query_execution(%{"QueryExecutionId" => exec_id})
|> elem(1)
|> then(& &1["QueryExecution"]["Status"]["State"])
```

![スクリーンショット 2022-12-05 15.12.46.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/84143591-1de5-dc4e-a30a-a6c1a7244883.png)

成功したようです

## SQL 実行関数定義

毎回書くのは面倒なので、 SQL 実行と実行状況確認の関数を作っておきます

```elixir
exec_query = fn query ->
  client
  |> AWS.Athena.start_query_execution(%{
    "QueryString" => query,
    "ClientRequestToken" => UUID.uuid1(),
    "ResultConfiguration" => %{
      "OutputLocation" => "s3://#{bucket_name}"
    }
  })
  |> elem(1)
  |> then(& &1["QueryExecutionId"])
end
```

```elixir
get_state = fn exec_id ->
  client
  |> AWS.Athena.get_query_execution(%{"QueryExecutionId" => exec_id})
  |> elem(1)
  |> then(& &1["QueryExecution"]["Status"]["State"])
end
```

## テーブル作成

次にテーブルを作成します

CSV からよしなに作ってほしいところですが、自分で定義してあげないといけません

再実行時にエラーにならないよう、もしテーブルがあれば削除しておきます

```elixir
exec_id = exec_query.("DROP TABLE IF EXISTS athena_sample.wine_table")
```

```elixir
get_state.(exec_id)
```

![スクリーンショット 2022-12-05 18.00.20.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/5f95d7ef-14c2-e3dd-0b07-52d44eea8f68.png)

CREATE 文の列名と型をいちいち書きたくないので、データフレームから生成します

列名とタイプのタプル配列を生成します

（列の順序が合っていないといけないため、 names で取得する順で並ぶようにしています）

```elixir
types =
  wine_df
  |> DataFrame.names
  |> Enum.map(fn name ->
    {name, DataFrame.dtypes(wine_df)[name]}
  end)
```

![スクリーンショット 2022-12-05 17.59.03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/466852df-6c27-e592-70a3-e3f670b6b8e8.png)

```elixir
# CREATE 文の列定義を生成
table_columns =
  types
  |> Enum.map(fn {name, type} ->
    name <>
      " " <>
      case type do
        :integer ->
          "int"

        :float ->
          "float"
      end
  end)
  |> Enum.join(",")
```

![スクリーンショット 2022-12-05 18.01.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/8bfef30b-6a17-aa64-b70f-d073e84c0c7a.png)

CREATE 文を定義して実行します

詳細についてはこちらの記事をご参照ください

https://qiita.com/yoshiyama_hana/items/3d532c7ecc5f08c0d040

```elixir
exec_id = exec_query.("
  CREATE EXTERNAL TABLE
    athena_sample.wine_table (#{table_columns})
  ROW FORMAT
    SerDe 'org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe' 
  WITH SerDeProperties (
    'serialization.format' = ',',
    'field.delim' = ','
  )
  STORED AS TEXTFILE
  LOCATION
    's3://#{bucket_name}/#{athena_prefix}' 
  TBLPROPERTIES (
    'has_encrypted_data'='false',
    'skip.header.line.count'='1'
  )
  ")
```

CSV　の場合、 `ROW FORMAT ...` 以降を指定しておけばおおよそ問題ありません

`LOCATION` で CSV をアップロードしたバケット、パスを指定します

```elixir
get_state.(exec_id)
```

![スクリーンショット 2022-12-05 15.22.44.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/6622bee2-be52-aa5e-f641-4e1acf22d769.png)

## データ取得

いよいよ SQL でデータを取得してみます

まずは全件全項目です

```elixir
exec_id = exec_query.("
  SELECT
    *
  FROM
    athena_sample.wine_table
  ")
```

```elixir
get_state.(exec_id)
```

![スクリーンショット 2022-12-05 15.26.07.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/78fd97f6-7625-2ae8-e5b1-2a7a82ab38e1.png)

とりあえず成功はしましたが、結果はまた別途取得しないといけません

```elixir
results =
  client
  |> AWS.Athena.get_query_results(%{"QueryExecutionId" => exec_id})
  |> elem(1)
```

![スクリーンショット 2022-12-05 18.03.07.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/e0ddfb29-ec9b-1d37-c026-f536fbc634a4.png)

何やら色々返ってきていますが、 ResultSet の中の ResultSetMetadata に結果の列情報が入っています

ここから各列の名前と型を取得します

```elixir
# 列名と型を取得
types =
  results
  |> then(& &1["ResultSet"]["ResultSetMetadata"]["ColumnInfo"])
  |> Enum.map(fn info ->
    {
      info["Name"],
      case info["Type"] do
        "integer" ->
          :integer

        "float" ->
          :float

        _ ->
          :string
      end
    }
  end)
```

![スクリーンショット 2022-12-05 18.03.53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/d2b0a231-95fc-7389-49d7-a835253de971.png)

また、 ResultSet の Rows に結果の各行が入っています

ただし、先頭行はヘッダーなので無視します

扱いやすいようにデータフレームにしておきましょう

```elixir
# データを取得
results_df =
  results
  |> then(& &1["ResultSet"]["Rows"])
  |> Enum.map(& &1["Data"])
  |> then(fn columns ->
    [headers | values] = columns

    headers
    |> Enum.map(& &1["VarCharValue"])
    |> Enum.with_index()
    |> Enum.reduce(%{}, fn {col_name, index}, acc ->
      col_values =
        Enum.map(values, fn each_values ->
          each_values
          |> Enum.at(index)
          |> then(& &1["VarCharValue"])
        end)

      %{col_name => col_values}
      |> Map.merge(acc)
    end)
  end)
  |> DataFrame.new()
```

![スクリーンショット 2022-12-05 18.04.55.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/11b511f5-9129-f521-2e75-ce42cfbc88ef.png)

取得してきた時点では全列が文字列型になっているので、型情報を元に変換します

```elixir
parsed_df =
  types
  |> Enum.reduce(results_df, fn {col_name, col_type}, df ->
    DataFrame.mutate_with(df, &%{col_name => Series.cast(&1[col_name], col_type)})
  end)
```

![スクリーンショット 2022-12-05 18.04.24.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/2f919ffe-7f59-8d44-6513-ddd05a79e274.png)

データテーブルで表示してみましょう

```elixir
parsed_df
|> Kino.DataTable.new()
```

![スクリーンショット 2022-12-05 18.05.30.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/b0e83eda-6284-5766-1f96-ac45b9a4d99c.png)

アップロードした元データと同じく、 178 件取得できました

データ取得も何かと面倒なので、関数化しておきます

```elixir
get_results = fn exec_id ->
  results =
    client
    |> AWS.Athena.get_query_results(%{"QueryExecutionId" => exec_id})
    |> elem(1)

  results_df =
    results
    |> then(& &1["ResultSet"]["Rows"])
    |> Enum.map(& &1["Data"])
    |> then(fn columns ->
      [headers | values] = columns

      headers
      |> Enum.map(& &1["VarCharValue"])
      |> Enum.with_index()
      |> Enum.reduce(%{}, fn {col_name, index}, acc ->
        col_values =
          Enum.map(values, fn each_values ->
            each_values
            |> Enum.at(index)
            |> then(& &1["VarCharValue"])
          end)

        %{col_name => col_values}
        |> Map.merge(acc)
      end)
    end)
    |> DataFrame.new()

  results
  |> then(& &1["ResultSet"]["ResultSetMetadata"]["ColumnInfo"])
  |> Enum.map(fn info ->
    {
      info["Name"],
      case info["Type"] do
        "integer" ->
          :integer

        "float" ->
          :float

        _ ->
          :string
      end
    }
  end)
  |> Enum.reduce(results_df, fn {col_name, col_type}, df ->
    DataFrame.mutate_with(df, &%{col_name => Series.cast(&1[col_name], col_type)})
  end)
end
```

次は `WHERE` や `ORDER BY` を使ってみましょう

```elixir
exec_id = exec_query.("
  SELECT
    class,
    color_intensity,
    flavanoids
  FROM
    athena_sample.wine_table
  WHERE
    color_intensity < 5.0
    AND flavanoids >= 2.0
  ORDER BY
    class
  ")
 ```

```elixir
exec_id
|> get_results.()
|> Kino.DataTable.new()
```

![スクリーンショット 2022-12-05 18.07.06.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/929fed36-e977-ad02-9f88-83d385360825.png)

見事に抽出、ソートができているようです

`GROUP BY` も使ってみましょう

```elixir
exec_id = exec_query.("
  SELECT
    class,
    max(alcohol) AS alcohol
  FROM
    athena_sample.wine_table
  GROUP BY
    class
  ORDER BY
    alcohol DESC
  ")
```

```elixir
exec_id
|> get_results.()
|> Kino.DataTable.new()
```

![スクリーンショット 2022-12-05 18.07.49.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/2c86cbbc-5d5e-6e35-3633-37e1f2c9442b.png)

しっかり集計できていますね

## テーブル削除

テーブルを消しておきます

```elixir
exec_id = exec_query.("DROP TABLE wine_table")
```

```elixir
get_state.(exec_id)
```

![スクリーンショット 2022-12-05 15.37.45.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/f9629b53-65f7-232f-855d-f40a833a7369.png)

## データベース削除

データベースを消しておきます

```elixir
exec_id = exec_query.("DROP DATABASE athena_sample")
```

```elixir
get_state.(exec_id)
```

![スクリーンショット 2022-12-05 15.37.45.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/f9629b53-65f7-232f-855d-f40a833a7369.png)

## DataFrame で同じ動作

SQL と同じことを DataFrame でそのままやってみましょう

```elixir
wine_df
|> DataFrame.filter(color_intensity < 5)
|> DataFrame.filter(flavanoids >= 2.0)
|> DataFrame.arrange(class)
|> DataFrame.select(["class", "color_intensity", "flavanoids"])
|> Kino.DataTable.new()
 ```

![スクリーンショット 2022-12-05 18.09.57.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/485179eb-dd99-ab34-9293-88c56bfb39c2.png)

```elixir
wine_df
|> DataFrame.group_by("class")
|> DataFrame.summarise(alcohol: max(alcohol))
|> DataFrame.arrange(desc: alcohol)
|> Kino.DataTable.new()
```

![スクリーンショット 2022-12-05 18.11.52.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/3cd074c8-57ba-c9d7-7230-17ebc4bd32b6.png)

当然できますね

## まとめ

というわけで、よほど大きいデータでローカルで開けない、というような事情でもない限り、 Athena は使わないと思います

DataFrame　なら SQL 以上に複雑な操作も可能ですしね

（ここまで書いておいて）