"""このサンプルは、Spark Streamingを使用してCOSストレージ内の推論テーブルをストリーミング読み取りし、主要フィールド（特徴、予測値）を加工処理した後、DLC内部テーブルに書き込む処理プロセスを示しています。"""

## 前提条件
1. このDLCエンジンに関連付けられたCOSアクセス認証情報に、推論テーブルストレージバケットへのアクセス権限があることを確認してください。
2. モデルサービスが推論テーブル監視を有効にしていることを確認してください。

## 実行方法
チュートリアルのパラメータを変更した後、オフライン開発->オーケストレーションスペースでDLC PySparkタスクを選択して実行します。

## 定時スケジュール

対応するパラメータを変更してデバッグを完了した後、このタスクを定時実行に設定し、推論テーブルデータを継続的に加工できます。また、加工後の推論テーブルデータをTagデータと関連付け、Tagデータを含む推論テーブルを生成し、データ品質モジュールで推論テーブルに機械学習関連の監視指標を設定して、モデル品質を継続的に監視できます。
"""
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql import functions as F
from pyspark.sql import types as T
from typing import Union, List, Dict, Optional, Sequence, Any
from pyspark.sql.streaming import StreamingQuery

from pyspark.sql.types import (
    StructType, StructField, StringType, LongType, DoubleType,
    ArrayType, TimestampType
)

"""
### 元の推論テーブルデータ形式の定義
元の推論テーブルデータは特定の構造に従ってJSON形式でユーザー指定のCOSバケットに保存されます。request_schemaはユーザー入力リクエスト構造を表し、response_schemaはモデル出力構造を表します。モデルサービスに対応する実際の入出力形式に応じて調整してください。その他のフィールドは推論テーブルの固定構造であるため、通常は調整不要です。
"""
request_schema = StructType([
    StructField("client_request_id", StringType(), True),
    StructField("dataframe_split", StructType([
        StructField("columns", ArrayType(StringType()), True),
        StructField("data", ArrayType(ArrayType(DoubleType())), True),
    ]), True),
])

request_metadata_schema = StructType([
    StructField("model_name", StringType(), True),
    StructField("model_version", StringType(), True),
    StructField("service_id", StringType(), True),
])

response_schema = StructType([
    StructField("predictions", ArrayType(LongType()), True),
])

schema = StructType([
    StructField("client_request_id", StringType(), True),
    StructField("execution_duration_ms", LongType(), True),
    StructField("request", request_schema, True),
    StructField("request_date", StringType(), True),
    StructField("request_metadata", request_metadata_schema, True),
    StructField("request_time", TimestampType(), True),
    StructField("requester", StringType(), True),
    StructField("response", response_schema, True),
    StructField("sampling_fraction", DoubleType(), True),
    StructField("status_code", LongType(), True),
    StructField("wedata_request_id", StringType(), True),
])

# データソースパラメータ
SOURCE_PATH = "cosn://ストレージバケットパス/"  # 推論テーブルストレージパス
CHECKPOINT_PATH = "cosn://ストレージバケットパス/checkpoint" # check_pointストレージパス、推論テーブルストレージパスと一致させることを推奨します

# 結果テーブルパラメータ
UNPACKED_TABLE_NAME = "DataLakeCatalog.test.unpacked_test_inference_table"  # 結果テーブル名。DLC内部テーブルの3セグメント形式のテーブル名である必要があります：<データカタログ>.<データベース>.<データテーブル>
MODEL_ID_COL = "model_id" # 結果テーブルモデルID列、モデルの識別に使用します
EXAMPLE_ID_COL = "example_id" # 結果テーブルレコードの一意のID列、ユーザー更新操作
PREDICTION_COL = "prediction" # 予測値列名
FEATURE_COLUMNS = ["sepal_length", "sepal_width", "petal_length", "petal_width"] # 特徴列名。モデルの特徴に応じて調整します


def process_requests(requests_raw: DataFrame) -> DataFrame:
    """
    リクエスト特徴を個別の列に展開し、対応する予測結果とペアリングします。

    Args:
        requests_raw: 処理対象データ
    Return:
        処理済みデータ
    """
    # 成功したリクエストをフィルタリングします
    requests_success = requests_raw.filter(F.col("status_code") == 200).drop("status_code")

    # モデル識別子を生成します
    requests_identified = requests_success \\
        .withColumn(MODEL_ID_COL, F.concat(F.col("request_metadata").getItem("model_name"), F.lit("_"), F.col("request_metadata").getItem("model_version"))) \\
        .drop("request_metadata")

    # 特徴列と予測結果を展開します
    # 1. 特徴列名、特徴データ、予測値を取得します
    requests_with_features = (
        requests_identified
        .withColumn("feature_columns", F.col("request.dataframe_split.columns"))
        .withColumn("feature_data", F.col("request.dataframe_split.data"))
        .withColumn("predictions", F.col("response.predictions"))
    )

    # 2. 各データ行を対応する予測結果とペアリングします。
    requests_exploded = (
        requests_with_features
        .withColumn("feature_prediction_pairs", 
                   F.arrays_zip(F.col("feature_data"), F.col("predictions")))
        .withColumn("feature_prediction_pairs", 
                   F.explode(F.col("feature_prediction_pairs")))
        .withColumn("feature_row", F.col("feature_prediction_pairs.feature_data"))
        .withColumn(PREDICTION_COL, F.col("feature_prediction_pairs.predictions"))
    )

    # 3. 特徴列を動的に作成します
    requests_with_feature_cols = requests_exploded
    for i, col_name in enumerate(FEATURE_COLUMNS):
        requests_with_feature_cols = (
            requests_with_feature_cols
            .withColumn(col_name, F.col("feature_row")[i])
        )
    
    # 4. レコードの一意のIDを追加し、後続のupsert操作を容易にします。
    requests_with_example_id = requests_with_feature_cols \\
        .withColumn(EXAMPLE_ID_COL, F.expr("uuid()"))

    # 5. 一時列をクリーンアップします。
    requests_processed = (
        requests_with_example_id
        .drop("feature_columns", "feature_data", "feature_prediction_pairs", "feature_row", "predictions", "request", "response")
    )

    return requests_processed


def create_table(
        spark: SparkSession,
        table_name: str,
        df: Optional[DataFrame] = None,
        partition_expr: str = None,
        description: Optional[str] = None,
):

    """
    dlc内部テーブルを作成します。

    Args:
        spark：SparkSessionインスタンス
        table_name: テーブル完全名（形式：データカタログ.データベース.データテーブル）
        df: 初期データ（スキーマ推論に使用）
        partition_expr: パーティション列（ストレージクエリを最適化）
        description: テーブル説明
    Raises:
        ValueError：異常作成時にスローされます。
    """
    # テーブルスキーマを推論します。
    table_schema = df.schema
    
    # 列定義を構築します。
    columns_ddl = []
    for field in table_schema.fields:
        data_type = field.dataType.simpleString().upper()
        col_def = f"`{field.name}` {data_type}"
        if not field.nullable:
            col_def += " NOT NULL"
        columns_ddl.append(col_def)
    
    # 建表文を構築します。
    ddl = f"""
    CREATE TABLE IF NOT EXISTS {table_name} (
        {', '.join(columns_ddl)}
    )
    USING iceberg
    PARTITIONED BY ({partition_expr})
    TBLPROPERTIES (
        'comment'= '{description or ''}',
        'format-version'= '2',
        'write.metadata.previous-versions-max'= '100',
        'write.metadata.delete-after-commit.enabled'= 'true',
        'smart-optimizer.inherit' = 'none',
        'smart-optimizer.written.enable' = 'enable'
    )
    """

    # sqlを出力します。
    print(f"create table ddl: {ddl}\\n")

    # DDLを実行します。
    try:
        spark.sql(ddl)
    except Exception as e:
        raise ValueError(f"Failed to create table: {str(e)}") from e

    print(f"create table {table_name} done")

spark = SparkSession.builder.appName("Operate DB Example").getOrCreate()

# ストリーミングDataFrameを構築します。
df = (
    spark.readStream.format("json")
    .schema(schema)
    .load(SOURCE_PATH)
)

# ストリーミングデータを処理します。
unpacked_df = process_requests(df)

# 結果テーブルを作成し、request_timeをパーティションフィールドとして、日ごとにパーティションを作成します。モデルサービスのリクエスト量に応じて調整できます。
create_table(spark, UNPACKED_TABLE_NAME, unpacked_df, 'days(request_time)')

# データを書き込みます。
query = unpacked_df.writeStream \\
        .format("parquet") \\
        .outputMode("append") \\
        .option("checkpointLocation", CHECKPOINT_PATH) \\
        .trigger(once=True) \\
        .toTable(UNPACKED_TABLE_NAME)

query.awaitTermination()  


spark.sql(f"select * from {UNPACKED_TABLE_NAME}").show(10)



"""
このサンプルは、Spark Streamingを使用してCOSストレージ内の推論テーブルをストリーミング読み取りし、主要フィールド（特徴、予測値）を加工処理した後、EMR Hiveテーブルに書き込む処理プロセスを示しています。

## 前提条件
1. このDLCエンジンに関連付けられたCOSアクセス認証情報に、推論テーブルストレージバケットへのアクセス権限があることを確認してください。
2. モデルサービスが推論テーブル監視を有効にしていることを確認してください。

## 実行方法
チュートリアルのパラメータを変更した後、オフライン開発->オーケストレーションスペースでEMR PySparkタスクを選択して実行します。

## 定時スケジュール

対応するパラメータを変更してデバッグを完了した後、このタスクを定時実行に設定し、推論テーブルデータを継続的に加工できます。また、加工後の推論テーブルデータをTagデータと関連付け、Tagデータを含む推論テーブルを生成し、データ品質モジュールで推論テーブルに機械学習関連の監視指標を設定して、モデル品質を継続的に監視できます。

"""
from pyspark.sql import DataFrame, SparkSession
from pyspark.sql import functions as F
from pyspark.sql import types as T
from typing import Union, List, Dict, Optional, Sequence, Any
from pyspark.sql.streaming import StreamingQuery

from pyspark.sql.types import (
    StructType, StructField, StringType, LongType, DoubleType,
    ArrayType, TimestampType
)

"""
### 元の推論テーブルデータ形式の定義
元の推論テーブルデータは特定の構造に従ってJSON形式でユーザー指定のCOSバケットに保存されます。request_schemaはユーザー入力リクエスト構造を表し、response_schemaはモデル出力構造を表します。モデルサービスに対応する実際の入出力形式に応じて調整してください。その他のフィールドは推論テーブルの固定構造であるため、通常は調整不要です。
"""

request_schema = StructType([
    StructField("client_request_id", StringType(), True),
    StructField("dataframe_split", StructType([
        StructField("columns", ArrayType(StringType()), True),
        StructField("data", ArrayType(ArrayType(DoubleType())), True),
    ]), True),
])

request_metadata_schema = StructType([
    StructField("model_name", StringType(), True),
    StructField("model_version", StringType(), True),
    StructField("service_id", StringType(), True),
])

response_schema = StructType([
    StructField("predictions", ArrayType(LongType()), True),
])

schema = StructType([
    StructField("client_request_id", StringType(), True),
    StructField("execution_duration_ms", LongType(), True),
    StructField("request", request_schema, True),
    StructField("request_date", StringType(), True),
    StructField("request_metadata", request_metadata_schema, True),
    StructField("request_time", TimestampType(), True),
    StructField("requester", StringType(), True),
    StructField("response", response_schema, True),
    StructField("sampling_fraction", DoubleType(), True),
    StructField("status_code", LongType(), True),
    StructField("wedata_request_id", StringType(), True),
])


# データソースパラメータ
SOURCE_PATH = "cosn://ストレージバケットパス/"  # 推論テーブルストレージパス
CHECKPOINT_PATH = "cosn://ストレージバケットパス/checkpoint" # check_pointストレージパス、推論テーブルストレージパスと一致させることを推奨します

# 結果テーブルパラメータ
UNPACKED_TABLE_NAME = "testdb.unpacked_test_inference_table"  # 結果テーブル名。hiveテーブル名である必要があります：<データベース>.<データテーブル>
MODEL_ID_COL = "model_id" # 結果テーブルモデルID列、モデルの識別に使用します
EXAMPLE_ID_COL = "example_id" # 結果テーブルレコードの一意のID列、ユーザー更新操作
PREDICTION_COL = "prediction" # 予測値列名
FEATURE_COLUMNS = ["sepal_length", "sepal_width", "petal_length", "petal_width"] # 特徴列名。モデルの特徴に応じて調整します

def process_requests(requests_raw: DataFrame) -> DataFrame:
    """
    リクエスト特徴を個別の列に展開し、対応する予測結果とペアリングします。

    Args:
        requests_raw: 処理対象データ
    Return:
        処理済みデータ
    """
    # 成功したリクエストをフィルタリングします
    requests_success = requests_raw.filter(F.col("status_code") == 200).drop("status_code")

    # モデル識別子を生成します
    requests_identified = requests_success \\
        .withColumn(MODEL_ID_COL, F.concat(F.col("request_metadata").getItem("model_name"), F.lit("_"), F.col("request_metadata").getItem("model_version"))) \\
        .drop("request_metadata")

    # 特徴列と予測結果を展開します
    # 1. 特徴列名、特徴データ、予測値を取得します
    requests_with_features = (
        requests_identified
        .withColumn("feature_columns", F.col("request.dataframe_split.columns"))
        .withColumn("feature_data", F.col("request.dataframe_split.data"))
        .withColumn("predictions", F.col("response.predictions"))
    )

    # 2. 各データ行を対応する予測結果とペアリングします。
    requests_exploded = (
        requests_with_features
        .withColumn("feature_prediction_pairs", 
                   F.arrays_zip(F.col("feature_data"), F.col("predictions")))
        .withColumn("feature_prediction_pairs", 
                   F.explode(F.col("feature_prediction_pairs")))
        .withColumn("feature_row", F.col("feature_prediction_pairs.feature_data"))
        .withColumn(PREDICTION_COL, F.col("feature_prediction_pairs.predictions"))
    )

    # 3. 特徴列を動的に作成します
    requests_with_feature_cols = requests_exploded
    for i, col_name in enumerate(FEATURE_COLUMNS):
        requests_with_feature_cols = (
            requests_with_feature_cols
            .withColumn(col_name, F.col("feature_row")[i])
        )
    
    # 4. レコードの一意のIDを追加し、後続のupsert操作を容易にします。
    requests_with_example_id = requests_with_feature_cols \\
        .withColumn(EXAMPLE_ID_COL, F.expr("uuid()"))

    # 5. 一時列をクリーンアップします。
    requests_processed = (
        requests_with_example_id
        .drop("feature_columns", "feature_data", "feature_prediction_pairs", "feature_row", "predictions", "request", "response")
    )

    return requests_processed


def create_table(
        spark: SparkSession,
        table_name: str,
        df: Optional[DataFrame] = None,
        partition_expr: str = None,
        description: Optional[str] = None,
):
    """
    hiveテーブルを作成します。

    Args:
        spark：SparkSessionインスタンス
        table_name: hiveテーブル完全名（形式：データベース.データテーブル）
        df: 初期データ（オプション、スキーマ推論に使用）
        partition_expr: パーティション列（ストレージクエリの最適化）
        description: テーブル説明
    Raises:
        ValueError：異常作成時にスローされます。
    """
    # テーブルスキーマを推論します。
    table_schema = df.schema

    part_col_nam = partition_expr.split(' ')[0]
    
    # 列定義を構築します。
    columns_ddl = []
    for field in table_schema.fields:
        # テーブル作成時にパーティション列を除外します。
        if field.name == part_col_nam:
            continue
        data_type = field.dataType.simpleString().upper()
        col_def = f"`{field.name}` {data_type}"
        if not field.nullable:
            col_def += " NOT NULL"
        columns_ddl.append(col_def)
    
    # 建表文を構築します。
    ddl = f"""
    CREATE TABLE IF NOT EXISTS {table_name} (
        {', '.join(columns_ddl)}
    )
    PARTITIONED BY ({partition_expr})
    TBLPROPERTIES (
        'comment'= '{description or ''}'
    )
    """
    
    # sqlを出力します。
    print(f"create table ddl: {ddl}\\n")

    # DDLを実行します。
    try:
        spark.sql(ddl)
    except Exception as e:
        raise ValueError(f"Failed to create table: {str(e)}") from e

    print(f"create table {table_name} done")


def write_batch_to_hive(batch_df: DataFrame, batch_id: int, table_name: str):
    # 小ファイルの制御：パーティション列による再パーティション
    out_df = batch_df.repartition(8, "dt")
    
    out_df.write.mode("append").format("hive").saveAsTable(
        table_name,
        partitionBy="dt"
    )

spark = SparkSession.builder.appName("Operate DB Example").enableHiveSupport().getOrCreate()

# ストリーミングDataFrameを構築します。
df = (
    spark.readStream.format("json")
    .schema(schema)
    .load(SOURCE_PATH)
)

# ストリーミングデータを処理します。
unpacked_df = process_requests(df)

# パーティション列を挿入します。
requests_with_dt = (
        unpacked_df
        .withColumn("dt", F.substring(F.col("request_date"), 1, 10))  # YYYY-MM-DD形式の日付を抽出します。
)

# 結果テーブルを作成します。
create_table(spark, UNPACKED_TABLE_NAME, unpacked_df, 'dt STRING')

# Hive動的パーティションを設定します。
spark.sql("SET hive.exec.dynamic.partition=true")
spark.sql("SET hive.exec.dynamic.partition.mode=nonstrict")

# 結果テーブルにバッチで書き込みます。
query = (
    requests_with_dt
    .writeStream
    .outputMode("append")
    .option("checkpointLocation", CHECKPOINT_PATH)
    .trigger(once=True)
    .foreachBatch(lambda df, id: write_batch_to_hive(df, id, UNPACKED_TABLE_NAME))
    .start()
)

query.awaitTermination()

# 結果データを照会します。
spark.sql(f"select * from {UNPACKED_TABLE_NAME}").show(10)