【S3 Tables】メンテナンスジョブのファイル圧縮について調べてみよう

import time
import boto3
import random
import string
import pyarrow as pa
import datetime as dt
from pyiceberg.catalog import load_catalog

# S3 Tablesカタログの読み込み関数
def init_rest_catalog(
    s3table_bucket_name, account_id, region, catalog_name="S3TablesCatalog"
):
    return load_catalog(
        catalog_name,
        **{
            "type": "rest",
            "warehouse": f"{account_id}:s3tablescatalog/{s3table_bucket_name}",
            "uri": f"https://glue.{region}.amazonaws.com/iceberg",
            "rest.sigv4-enabled": "true",
            "rest.signing-name": "glue",
            "rest.signing-region": region,
        },
    )

# ランダムな文字列を生成
def randstr(n):
    randlst = [random.choice(string.ascii_letters) for i in range(n)]
    return "".join(randlst)

def lambda_handler(event, context):

    # 現在の日本時刻を取得
    td = dt.timedelta(hours=9)
    now_jst = dt.datetime.now() + td

    # ================================================
    # カタログ読み込み用の変数
    region = "ap-northeast-1"
    account_id = "123456789012"
    s3table_bucket_name = "table-bucket-name"

    # テーブル作成先のネームスペースとテーブル名
    database_name = "name_space_name"
    table_name = "tbl_physcl_dlt"
    mttable_name = "tbl_physcl_dlt_mtdt"
    # ================================================

    # カタログを読み込み
    catalog = init_rest_catalog(s3table_bucket_name, account_id, region)

    # テーブルを読み込み
    tbl = catalog.load_table(f"{database_name}.{table_name}")
    mttbl = catalog.load_table(f"{database_name}.{mttable_name}")

    # ================================================
    # ダミーデータの投入と削除
    # ================================================

    # ランダム文字列と乱数で1000行分のデータを作成
    dm_dict = {"col_1": [], "col_2": [], "col_3": []}
    for i in range(1000):

        dm_dict["col_1"].append(randstr(12))
        dm_dict["col_2"].append(random.randint(1, 999))
        dm_dict["col_3"].append(random.choice([True, False]))

    # データをテーブルに変換してINSERT
    schema = pa.schema(
        [
            pa.field("col_1", pa.string(), nullable=True),
            pa.field("col_2", pa.int32(), nullable=True),
            pa.field("col_3", pa.bool_(), nullable=True),
        ]
    )
    padf = pa.Table.from_pydict(dm_dict, schema=schema)
    tbl.append(padf)

    # col_2が400以下の行を削除する
    # tbl.delete(delete_filter="col_2 <= 400")
    # pyicebergから実行するとMoRは対応してなかったため、Athenaで実行
    athena_client = boto3.client("athena")
    athena_client.start_query_execution(
        QueryString=f"delete from tbl_physcl_dlt where col_2 <= 400;",
        QueryExecutionContext={
            "Catalog": f"s3tablescatalog/{s3table_bucket_name}",
            "Database": database_name,
        },
        ResultConfiguration={"OutputLocation": "s3://bucket_name"},
        WorkGroup="workgroup_name",
    )

    # Athenaのクエリ実行に時間がかかる可能性があるのでスリープ
    # 本当はポーリングした方がいい
    time.sleep(60)
    # テーブルの最新の状態を取得
    tbl = catalog.load_table(f"{database_name}.{table_name}")

    # ================================================
    # メタデータを取得と格納
    # ================================================

    # メタデータ用の辞書
    mt_dict = {}

    # 全ファイルのファイルサイズ合計を取得
    files = tbl.inspect.files().to_pydict()
    mt_dict["files_size"] = [sum(files["file_size_in_bytes"])]

    # データファイルのファイル数とファイルサイズ合計を取得
    data_files = tbl.inspect.data_files().to_pydict()
    mt_dict["data_files_count"] = [len(data_files["content"])]
    mt_dict["data_files_size"] = [sum(data_files["file_size_in_bytes"])]

    # 削除ファイルのファイル数とファイルサイズ合計を取得
    delete_files = tbl.inspect.delete_files().to_pydict()
    mt_dict["delete_files_count"] = [len(delete_files["content"])]
    mt_dict["delete_files_size"] = [sum(delete_files["file_size_in_bytes"])]

    # スナップショット数を取得
    mt_dict["snapshots_count"] = [len(tbl.snapshots())]

    # メンテンスジョブの履歴を取得
    s3tbl_client = boto3.client("s3tables")
    response = s3tbl_client.get_table_maintenance_job_status(
        tableBucketARN=f"arn:aws:s3tables:{region}:{account_id}:bucket/{s3table_bucket_name}",
        namespace=database_name,
        name=table_name,
    )

    # メンテナンスジョブのレスポンスからタイムスタンプを取得
    for k, v in response["status"].items():

        try:

            # 最終実行時日時を日本時間に変換
            ts_jst = v["lastRunTimestamp"] + td
            ts_jst_str = dt.datetime.strftime(ts_jst, "%Y-%m-%d %H:%M:%S")

        except:

            # 取得できなかった場合
            ts_jst_str = "Not_Yet_Run"

        # カラムと日時の値を格納
        mt_dict[k.lower()] = [ts_jst_str]

    # 実行時刻を格納
    mt_dict["timestamp"] = [now_jst.strftime("%Y-%m-%d %H:%M:%S")]

    # メタデータ情報を格納
    mt_schema = pa.schema(
        [
            pa.field("files_size", pa.int32(), nullable=True),
            pa.field("data_files_count", pa.int32(), nullable=True),
            pa.field("data_files_size", pa.int32(), nullable=True),
            pa.field("delete_files_count", pa.int32(), nullable=True),
            pa.field("delete_files_size", pa.int32(), nullable=True),
            pa.field("snapshots_count", pa.int32(), nullable=True),
            pa.field("icebergcompaction", pa.string(), nullable=True),
            pa.field("icebergunreferencedfileremoval", pa.string(), nullable=True),
            pa.field("icebergsnapshotmanagement", pa.string(), nullable=True),
            pa.field("timestamp", pa.string(), nullable=True),
        ]
    )
    mt_padf = pa.Table.from_pydict(mt_dict, schema=mt_schema)
    mttbl.append(mt_padf)

    return {"status_code": 200, "results": mt_dict}