LoginSignup
8

More than 3 years have passed since last update.

AWS SageMakerからathenaにクエリを出し、pandas DataFrameにデータを格納する

Posted at

3行で

・SageMakerのノートブックインスタンスからathenaにあるDBにアクセスし、データを作りたい
・pyathenaという便利なモジュールがある
https://pypi.org/project/PyAthena/

モチベーション

SageMakerのノートブックは環境構築の手間が軽く済んでとても便利。AWS内だから他のサービスとも連携しやすい。
通常、大きいデータはS3なりRedshiftにあると思われるので、SageMakerで学習を行うなら、SageMakerからそれらにアクセスする必要がある。
S3にあるデータを使って特徴量エンジニアリングをするには、特にメモリに乗るかどうかぐらいの大きなデータの場合、一旦athenaに取り込んでSQLベースでデータを作る方法も選択肢になる。
SageMakerのノートブックからathenaにクエリを出し、クエリ結果をpandasのDFとして受け取るにはどうしたらいいだろうか。

結論

まず、ノートブックインスタンスのIAM ロールにathenaへのアクセス権をあげましょう。
アクセス権をあげた上で、SageMakerのノートブックに書くべきサンプルコードは以下

from pyathena import connect
import pandas as pd
import boto3

#set parameter
session = boto3.Session(profile_name=None)
REGION = session.region_name
ATHENA_STAGING = 'your_bucket'

#make connection
conn = connect(s3_staging_dir=ATHENA_STAGING,
               region_name=REGION)
query = """
SELECT * 
FROM YourDatabase.YourTable 
LIMIT 8;
"""

#query
df = pd.read_sql(query, conn)

set parameterのブロックは、connectに渡すパラメータを設定するパート
別にこのやり方でなくても、REGIONなんてハードコーディングでもいいと思います。
your_bucketにはathenaのQuery result locationが入ります。
athenaのコンソールのsettings -> Query result locationにs3://~~~ のURIが入ってますんでそれです。

queryにはクエリするSQL文が入ります。
こんな感じで一度文字列変数に代入すると、コードが見やすくなります。
athenaのクエリエディタで文法レベルの書き間違いがないか確認してからSageMakerで使うとデバッグが速く進むでしょう。

用例

ここまでやると、SageMakerのノートブックから好き放題クエリ投げられます。
前の記事で紹介したhash値を使ったtrain/test splitもよし
特徴量選択のために、カテゴリカルな特徴量に集計関数を使って各カテゴリーごとの推論対象との相関をみるもよし。
それは例えばこんな感じのquery生成用関数を作るととても捗る。

def get_distinct_query(column_name):
  sql = """
SELECT
  {0},
  COUNT(1) AS count,
  AVG(target) AS avg
FROM
  YourDatabase.YourTable
WHERE
  MOD(ABS(FROM_BIG_ENDIAN_64(XXHASH64(TO_UTF8(<your_key>)))),10) < 1
GROUP BY
  {0}
ORDER BY {0} ASC
  """.format(column_name)
  return sql

ま、pandasだけでできる場合もあるけどね!
ビッグなデータにご用心。

参考文献

pyathena
https://pypi.org/project/PyAthena/

GCPを使った機械学習のcourseraの講義資料
https://github.com/GoogleCloudPlatform/training-data-analyst/blob/master/courses/machine_learning/deepdive/06_structured/1_explore.ipynb

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
8