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 = 'リージョン'
    )

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

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

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

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

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)

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)

import urllib
scheme, netloc, path, params, query, fragment = urllib.parse.urlparse(data_location)
print(netloc)
print(path)

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)

# テーブルを削除
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')