# =====================================================
# 本番用Sparkジョブ テンプレート
# =====================================================

import sys
import logging
from datetime import datetime
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# ロガーの設定
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def create_spark_session(app_name):
    """SparkSessionを作成"""
    return SparkSession.builder \
        .appName(app_name) \
        .getOrCreate()


def validate_input(spark, input_path):
    """入力データの検証"""
    try:
        df = spark.read.parquet(input_path)
        count = df.count()
        logger.info(f"入力データ件数: {count}")
        
        if count == 0:
            raise ValueError("入力データが空です")
        
        return df
    except Exception as e:
        logger.error(f"入力データの読み込みに失敗: {e}")
        raise


def process_data(df):
    """メインの処理ロジック"""
    logger.info("データ処理を開始...")
    
    # ここに実際の処理を記述
    result = df.groupBy("category").agg(
        F.sum("amount").alias("total_amount"),
        F.count("*").alias("count")
    )
    
    logger.info("データ処理完了")
    return result


def save_output(df, output_path):
    """結果の保存"""
    logger.info(f"出力先: {output_path}")
    
    df.write.parquet(output_path, mode="overwrite")
    
    logger.info("出力完了")


def main(input_path, output_path):
    """メイン処理"""
    start_time = datetime.now()
    logger.info(f"ジョブ開始: {start_time}")
    
    spark = None
    exit_code = 0
    
    try:
        # SparkSession作成
        spark = create_spark_session("MySparkJob")
        
        # 入力データ読み込み・検証
        df = validate_input(spark, input_path)
        
        # 処理実行
        result = process_data(df)
        
        # 結果保存
        save_output(result, output_path)
        
        logger.info("ジョブ正常終了")
        
    except Exception as e:
        logger.error(f"ジョブ失敗: {e}")
        exit_code = 1
        
    finally:
        if spark:
            spark.stop()
        
        end_time = datetime.now()
        duration = end_time - start_time
        logger.info(f"処理時間: {duration}")
    
    sys.exit(exit_code)


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: spark-submit job.py <input_path> <output_path>")
        sys.exit(1)
    
    input_path = sys.argv[1]
    output_path = sys.argv[2]
    
    main(input_path, output_path)

# ジョブの実行例

# ローカル実行
spark-submit job.py \
    /path/to/input \
    /path/to/output

# EMR / Dataprocでの実行
spark-submit job.py \
    s3://bucket/input \
    s3://bucket/output

# または
spark-submit job.py \
    gs://bucket/input \
    gs://bucket/output

# 適切なエラーハンドリング

def robust_data_load(spark, path, schema=None):
    """堅牢なデータ読み込み"""
    try:
        if schema:
            df = spark.read.schema(schema).parquet(path)
        else:
            df = spark.read.parquet(path)
        
        # 空データチェック
        if df.rdd.isEmpty():
            logger.warning(f"警告: データが空です - {path}")
        
        return df
        
    except AnalysisException as e:
        # ファイルが見つからない場合
        logger.error(f"ファイルが見つかりません: {path}")
        raise FileNotFoundError(f"入力ファイルが存在しません: {path}")
        
    except Exception as e:
        logger.error(f"データ読み込みエラー: {e}")
        raise

import time

def retry_operation(func, max_retries=3, delay=10):
    """リトライ機能付き関数実行"""
    for attempt in range(max_retries):
        try:
            return func()
        except Exception as e:
            if attempt < max_retries - 1:
                logger.warning(
                    f"処理失敗（{attempt + 1}/{max_retries}）: {e}"
                    f" {delay}秒後にリトライ..."
                )
                time.sleep(delay)
            else:
                logger.error(f"最大リトライ回数に達しました: {e}")
                raise

# 使用例
def save_with_retry(df, output_path):
    def _save():
        df.write.parquet(output_path, mode="overwrite")
    
    retry_operation(_save, max_retries=3, delay=30)

def validate_data_quality(df, checks):
    """データ品質チェック"""
    errors = []
    
    # NULL値チェック
    if "no_nulls" in checks:
        for col_name in checks["no_nulls"]:
            null_count = df.filter(F.col(col_name).isNull()).count()
            if null_count > 0:
                errors.append(f"{col_name}にNULL値が{null_count}件あります")
    
    # 必須カラムチェック
    if "required_columns" in checks:
        existing_cols = set(df.columns)
        for col_name in checks["required_columns"]:
            if col_name not in existing_cols:
                errors.append(f"必須カラム {col_name} がありません")
    
    # 重複チェック
    if "unique_columns" in checks:
        unique_cols = checks["unique_columns"]
        dup_count = df.groupBy(unique_cols).count() \
                      .filter(F.col("count") > 1).count()
        if dup_count > 0:
            errors.append(f"重複レコードが{dup_count}件あります")
    
    if errors:
        for error in errors:
            logger.error(error)
        raise ValueError("データ品質チェックに失敗しました")
    
    logger.info("データ品質チェックOK")
    return True

# 使用例
checks = {
    "no_nulls": ["user_id", "order_date"],
    "required_columns": ["user_id", "amount", "order_date"],
    "unique_columns": ["order_id"]
}
validate_data_quality(df, checks)

import json
from datetime import datetime

class StructuredLogger:
    """構造化ログを出力するクラス"""
    
    def __init__(self, job_name):
        self.job_name = job_name
        self.job_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    def _log(self, level, message, **kwargs):
        log_entry = {
            "timestamp": datetime.now().isoformat(),
            "level": level,
            "job_name": self.job_name,
            "job_id": self.job_id,
            "message": message,
            **kwargs
        }
        print(json.dumps(log_entry, ensure_ascii=False))
    
    def info(self, message, **kwargs):
        self._log("INFO", message, **kwargs)
    
    def error(self, message, **kwargs):
        self._log("ERROR", message, **kwargs)
    
    def metrics(self, **kwargs):
        """メトリクスを記録"""
        self._log("METRICS", "metrics_recorded", **kwargs)

# 使用例
logger = StructuredLogger("SalesAnalysisJob")

logger.info("処理開始", input_path="s3://bucket/input")
logger.info("データ読み込み完了", record_count=1000000)
logger.metrics(
    input_records=1000000,
    output_records=50000,
    duration_seconds=120
)

# config.yaml（設定ファイル）

development:
  input_path: /home/user/data/input
  output_path: /home/user/data/output
  spark:
    master: local[*]
    executor_memory: 2g

staging:
  input_path: s3://staging-bucket/input
  output_path: s3://staging-bucket/output
  spark:
    master: yarn
    executor_memory: 4g

production:
  input_path: s3://prod-bucket/input
  output_path: s3://prod-bucket/output
  spark:
    master: yarn
    executor_memory: 8g
    num_executors: 10

import yaml
import os

def load_config(env=None):
    """環境に応じた設定を読み込む"""
    env = env or os.getenv("SPARK_ENV", "development")
    
    config_path = os.path.join(
        os.path.dirname(__file__), 
        "config.yaml"
    )
    
    with open(config_path, 'r') as f:
        all_config = yaml.safe_load(f)
    
    if env not in all_config:
        raise ValueError(f"Unknown environment: {env}")
    
    return all_config[env]

# 使用例
config = load_config("production")
print(config["input_path"])  # s3://prod-bucket/input

import argparse

def parse_args():
    """コマンドライン引数をパース"""
    parser = argparse.ArgumentParser(description='Spark Job')
    
    parser.add_argument(
        '--env', 
        default='development',
        choices=['development', 'staging', 'production'],
        help='実行環境'
    )
    parser.add_argument(
        '--date',
        required=True,
        help='処理対象日 (YYYY-MM-DD)'
    )
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='ドライラン（保存しない）'
    )
    
    return parser.parse_args()

# 実行例：
# spark-submit job.py --env production --date 2024-01-15

# ❌ 悪い例：追記モード（再実行でデータ重複）
df.write.parquet(output_path, mode="append")

# ✅ 良い例：上書きモード（再実行しても同じ結果）
df.write.parquet(output_path, mode="overwrite")

# ✅ 良い例：日付パーティションで上書き
df.write.partitionBy("date").parquet(
    output_path, 
    mode="overwrite"
)

# =====================================================
# 本番用Sparkジョブ 完成版テンプレート
# =====================================================

import sys
import argparse
import logging
from datetime import datetime
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# ロガー設定
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s'
)
logger = logging.getLogger(__name__)


class SalesAnalysisJob:
    """売上分析ジョブ"""
    
    def __init__(self, args):
        self.args = args
        self.spark = None
        self.start_time = datetime.now()
        self.metrics = {}
    
    def run(self):
        """ジョブのメイン処理"""
        exit_code = 0
        
        try:
            logger.info(f"ジョブ開始: date={self.args.date}")
            
            # 1. SparkSession作成
            self.spark = self._create_spark_session()
            
            # 2. データ読み込み
            df = self._load_data()
            self.metrics['input_count'] = df.count()
            logger.info(f"入力件数: {self.metrics['input_count']}")
            
            # 3. データ処理
            result = self._process(df)
            self.metrics['output_count'] = result.count()
            
            # 4. 結果保存
            if not self.args.dry_run:
                self._save(result)
            else:
                logger.info("ドライラン: 保存をスキップ")
            
            logger.info("ジョブ正常終了")
            
        except Exception as e:
            logger.error(f"ジョブ失敗: {e}", exc_info=True)
            exit_code = 1
            
        finally:
            self._cleanup()
        
        return exit_code
    
    def _create_spark_session(self):
        return SparkSession.builder \
            .appName(f"SalesAnalysis_{self.args.date}") \
            .getOrCreate()
    
    def _load_data(self):
        input_path = f"{self.args.input_path}/date={self.args.date}"
        logger.info(f"読み込み: {input_path}")
        return self.spark.read.parquet(input_path)
    
    def _process(self, df):
        logger.info("データ処理開始...")
        result = df.groupBy("category").agg(
            F.sum("amount").alias("total_sales")
        )
        return result
    
    def _save(self, df):
        output_path = f"{self.args.output_path}/date={self.args.date}"
        logger.info(f"保存先: {output_path}")
        df.write.parquet(output_path, mode="overwrite")
    
    def _cleanup(self):
        if self.spark:
            self.spark.stop()
        
        duration = datetime.now() - self.start_time
        self.metrics['duration_seconds'] = duration.total_seconds()
        logger.info(f"処理時間: {duration}")
        logger.info(f"メトリクス: {self.metrics}")


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--date', required=True, help='処理日 (YYYY-MM-DD)')
    parser.add_argument('--input-path', required=True, help='入力パス')
    parser.add_argument('--output-path', required=True, help='出力パス')
    parser.add_argument('--dry-run', action='store_true', help='ドライラン')
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    job = SalesAnalysisJob(args)
    sys.exit(job.run())