PySpark Advanced

from pyspark.sql import SparkSession

# Créer une SparkSession (point d'entrée unique)
spark = SparkSession.builder \
    .appName("PySpark Advanced") \
    .config("spark.sql.adaptive.enabled", "true") \
    .getOrCreate()

print(f"Spark version: {spark.version}")
print(f"App name: {spark.sparkContext.appName}")

from pyspark.sql.functions import col, sum as spark_sum

# Créer des données de test
data = [(i, f"cat_{i % 5}", float(i * 10)) for i in range(1000)]
df = spark.createDataFrame(data, ["id", "category", "amount"])

# Pipeline avec transformations
result = (
    df
    .filter(col("amount") > 100)
    .select("category", "amount")
    .groupBy("category")
    .agg(spark_sum("amount").alias("total"))
)

# Voir le plan d'exécution
print("=" * 50)
print("PLAN D'EXÉCUTION (explain)")
print("=" * 50)
result.explain("formatted")

# Plan complet avec toutes les étapes
print("PLAN COMPLET (Parsed → Analyzed → Optimized → Physical)")
print("=" * 60)
result.explain(True)

# Pattern : arguments en ligne de commande (main.py)

# === Exemple de main.py pour spark-submit ===
example_main = '''
#!/usr/bin/env python3
"""Point d'entrée du job Spark."""

import argparse
from pyspark.sql import SparkSession

def parse_args():
    parser = argparse.ArgumentParser(description="ETL Pipeline")
    parser.add_argument("--date", required=True, help="Date de traitement (YYYY-MM-DD)")
    parser.add_argument("--env", default="dev", choices=["dev", "prod"])
    parser.add_argument("--input", required=True, help="Chemin input")
    parser.add_argument("--output", required=True, help="Chemin output")
    return parser.parse_args()

def main():
    args = parse_args()
    
    spark = SparkSession.builder \
        .appName(f"ETL-{args.env}-{args.date}") \
        .getOrCreate()
    
    # Charger les données
    df = spark.read.parquet(f"{args.input}/date={args.date}")
    
    # Transformations...
    result = df.filter(df.amount > 0)
    
    # Écrire
    result.write.mode("overwrite").parquet(f"{args.output}/date={args.date}")
    
    spark.stop()

if __name__ == "__main__":
    main()
'''

print(example_main)

from pyspark.sql.functions import spark_partition_id

# Créer un DataFrame
df = spark.range(0, 1000000)
print(f"Partitions initiales : {df.rdd.getNumPartitions()}")

# repartition = SHUFFLE (redistribue les données)
df_repart = df.repartition(10)
print(f"Après repartition(10) : {df_repart.rdd.getNumPartitions()}")

# coalesce = PAS DE SHUFFLE (combine les partitions existantes)
df_coal = df_repart.coalesce(3)
print(f"Après coalesce(3) : {df_coal.rdd.getNumPartitions()}")

# Voir la distribution des données par partition
print("\nDistribution après repartition(10):")
df_repart.groupBy(spark_partition_id().alias("partition")).count().orderBy("partition").show()

# Activer AQE (recommandé Spark 3.0+)
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

# Vérifier la configuration
print("AQE enabled:", spark.conf.get("spark.sql.adaptive.enabled"))
print("Skew Join enabled:", spark.conf.get("spark.sql.adaptive.skewJoin.enabled"))

from pyspark import StorageLevel
import time

# Créer un DataFrame avec calculs
df = spark.range(0, 5000000).withColumn("squared", col("id") ** 2)

# Sans cache : chaque action recalcule tout
start = time.time()
count1 = df.filter(col("squared") > 1000000).count()
count2 = df.filter(col("squared") < 100).count()
print(f"Sans cache : {time.time() - start:.2f}s")

# Avec cache : calcul une seule fois
df_cached = df.cache()

start = time.time()
count1 = df_cached.filter(col("squared") > 1000000).count()
count2 = df_cached.filter(col("squared") < 100).count()
print(f"Avec cache : {time.time() - start:.2f}s")

# ⚠️ IMPORTANT : libérer la mémoire !
df_cached.unpersist()
print("\n✅ Cache libéré avec unpersist()")

from pyspark.sql.functions import broadcast

# Grande table (10M lignes simulées)
big_df = spark.range(0, 1000000).withColumn("category_id", (col("id") % 100).cast("int"))

# Petite table (100 lignes)
small_df = spark.createDataFrame(
    [(i, f"Category {i}") for i in range(100)],
    ["category_id", "category_name"]
)

# ❌ SANS broadcast hint (Spark peut ou non broadcaster)
result_no_hint = big_df.join(small_df, "category_id")
print("Sans broadcast hint:")
result_no_hint.explain()

print("\n" + "="*50 + "\n")

# ✅ AVEC broadcast hint (force le broadcast)
result_broadcast = big_df.join(broadcast(small_df), "category_id")
print("Avec broadcast():")
result_broadcast.explain()

# Configuration pour gérer le skew automatiquement (AQE)
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "256MB")

print("AQE Skew Join activé")
print("Spark va automatiquement détecter et gérer les partitions skewed")

import os
import shutil

# Créer des données de test
test_data = spark.range(0, 100000).withColumn("category", (col("id") % 10).cast("string"))

# Sauvegarder en Parquet
output_path = "/tmp/test_parquet"
if os.path.exists(output_path):
    shutil.rmtree(output_path)
test_data.write.parquet(output_path)

# Lecture avec filtre → predicate pushdown
df_filtered = spark.read.parquet(output_path).filter(col("id") < 100)

print("Plan d'exécution avec Predicate Pushdown:")
df_filtered.explain()
# Observe "PushedFilters" dans le plan !

from pyspark.sql.functions import year, month, dayofmonth, lit
from datetime import datetime, timedelta
import random

# Créer des données avec dates
dates = [(datetime(2024, 1, 1) + timedelta(days=random.randint(0, 90))).strftime("%Y-%m-%d") 
         for _ in range(10000)]
data = [(i, dates[i % len(dates)], float(random.randint(10, 1000))) 
        for i in range(10000)]

df = spark.createDataFrame(data, ["id", "date", "amount"])
df = df.withColumn("date", col("date").cast("date"))
df = df.withColumn("year", year("date")).withColumn("month", month("date"))

# Écriture partitionnée
output_partitioned = "/tmp/partitioned_data"
if os.path.exists(output_partitioned):
    shutil.rmtree(output_partitioned)

df.write.partitionBy("year", "month").parquet(output_partitioned)

print("Structure créée:")
for root, dirs, files in os.walk(output_partitioned):
    level = root.replace(output_partitioned, '').count(os.sep)
    indent = ' ' * 2 * level
    print(f"{indent}{os.path.basename(root)}/")
    if level < 2:  # Limiter la profondeur
        for d in sorted(dirs)[:3]:
            print(f"{indent}  {d}/")

from pyspark.sql.types import StructType, StructField, LongType, StringType, DecimalType, DateType

# Définir un schéma explicite
schema = StructType([
    StructField("id", LongType(), nullable=False),
    StructField("customer_name", StringType(), nullable=True),
    StructField("amount", DecimalType(10, 2), nullable=True),  # Précision pour montants !
    StructField("transaction_date", DateType(), nullable=True)
])

print("Schéma défini:")
print(schema.simpleString())

# Utilisation
# df = spark.read.schema(schema).parquet("data/transactions/")

print("\n💡 Conseil : Utiliser DecimalType pour les montants financiers")
print("   Évite les erreurs d'arrondi des float/double")

from pyspark.sql.functions import udf, when, col
from pyspark.sql.types import StringType
import time

# Créer des données de test
df = spark.range(0, 500000).withColumn("amount", (col("id") % 2000).cast("double"))

# ❌ MAUVAIS : Python UDF
@udf(StringType())
def categorize_udf(amount):
    if amount > 1000:
        return "high"
    elif amount > 100:
        return "medium"
    return "low"

start = time.time()
result_udf = df.withColumn("category", categorize_udf(col("amount")))
result_udf.write.mode("overwrite").format("noop").save()  # Force l'exécution
udf_time = time.time() - start
print(f"❌ Python UDF : {udf_time:.2f}s")

# ✅ BON : Expression Spark native
start = time.time()
result_native = df.withColumn("category",
    when(col("amount") > 1000, "high")
    .when(col("amount") > 100, "medium")
    .otherwise("low")
)
result_native.write.mode("overwrite").format("noop").save()
native_time = time.time() - start
print(f"✅ Expression native : {native_time:.2f}s")

print(f"\n📊 Speedup : {udf_time/native_time:.1f}x plus rapide avec expression native !")

from pyspark.sql.functions import pandas_udf
import pandas as pd
import numpy as np

# Pandas UDF = vectorisé (traite des Series, pas des scalaires)
@pandas_udf("double")
def log_transform(s: pd.Series) -> pd.Series:
    return np.log1p(s)

# Test
df = spark.range(0, 100000).withColumn("value", col("id").cast("double"))

start = time.time()
result = df.withColumn("log_value", log_transform(col("value")))
result.write.mode("overwrite").format("noop").save()
print(f"Pandas UDF : {time.time() - start:.2f}s")

result.show(5)

# Configuration AQE complète (recommandée)
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.localShuffleReader.enabled", "true")

# Avec AQE, tu peux laisser shuffle.partitions élevé
# Spark optimisera automatiquement
spark.conf.set("spark.sql.shuffle.partitions", "1000")

print("✅ Configuration AQE optimale :")
print(f"   adaptive.enabled = {spark.conf.get('spark.sql.adaptive.enabled')}")
print(f"   coalescePartitions = {spark.conf.get('spark.sql.adaptive.coalescePartitions.enabled')}")
print(f"   skewJoin = {spark.conf.get('spark.sql.adaptive.skewJoin.enabled')}")
print(f"   shuffle.partitions = {spark.conf.get('spark.sql.shuffle.partitions')}")
print("\n💡 Avec AQE, Spark ajuste automatiquement le nombre de partitions !")

# URL de Spark UI pour cette session
print(f"🔍 Spark UI disponible sur : {spark.sparkContext.uiWebUrl}")
print("\n📊 Pour voir les métriques d'un job, lance une action puis consulte l'UI")

# Exemple : déclencher un job pour voir dans l'UI
df = spark.range(0, 1000000)
result = df.groupBy((col("id") % 100).alias("group")).count()
result.collect()  # Action qui déclenche le job

print("\n✅ Job exécuté - consulte Spark UI pour voir les stages et métriques")

import os
import shutil
from pyspark.sql.functions import *
from datetime import datetime, timedelta
import random
import time

# Setup : créer les données de test
print("📦 Création des données de test...")

# Configuration optimale
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

# Nettoyer les répertoires
for path in ["/tmp/transactions", "/tmp/products", "/tmp/customers", "/tmp/output"]:
    if os.path.exists(path):
        shutil.rmtree(path)

# 1. Transactions (5M lignes)
transactions = spark.range(0, 500000).select(
    col("id").alias("transaction_id"),
    (col("id") % 10000).alias("product_id"),
    (col("id") % 50000).alias("customer_id"),
    (rand() * 1000).alias("amount"),
    date_add(lit("2024-01-01"), (rand() * 90).cast("int")).alias("date")
)
transactions.write.parquet("/tmp/transactions")
print(f"Transactions : {transactions.count()} lignes")

# 2. Produits (10K lignes) - CSV intentionnellement
products_data = [(i, f"Product_{i}", f"Category_{i % 50}", float(10 + i % 100)) 
                 for i in range(10000)]
products = spark.createDataFrame(products_data, 
    ["product_id", "product_name", "category", "base_price"])
products.write.mode("overwrite").option("header", True).csv("/tmp/products")
print(f"Produits : {products.count()} lignes (CSV)")

# 3. Clients (500K lignes)
customers = spark.range(0, 50000).select(
    col("id").alias("customer_id"),
    concat(lit("Customer_"), col("id")).alias("customer_name"),
    (col("id") % 5).alias("segment")
)
customers.write.parquet("/tmp/customers")
print(f"Clients : {customers.count()} lignes")

print("\n Données créées avec succès !")

# ❌ VERSION NON OPTIMISÉE (baseline)
print("="*50)
print("❌ PIPELINE NON OPTIMISÉ")
print("="*50)

# Désactiver AQE pour le baseline
spark.conf.set("spark.sql.adaptive.enabled", "false")

start = time.time()

# Lecture
transactions_df = spark.read.parquet("/tmp/transactions")
products_df = spark.read.option("header", True).option("inferSchema", True).csv("/tmp/products")  # ❌ Inférence
customers_df = spark.read.parquet("/tmp/customers")

# UDF non optimisé
@udf(StringType())
def get_amount_category(amount):
    if amount > 500: return "high"
    elif amount > 100: return "medium"
    return "low"

# Joins sans broadcast
result = transactions_df \
    .join(products_df, "product_id") \
    .join(customers_df, "customer_id") \
    .withColumn("amount_category", get_amount_category(col("amount"))) \
    .groupBy("category", "segment", "amount_category") \
    .agg(
        count("*").alias("num_transactions"),
        sum("amount").alias("total_amount")
    )

# Écriture non partitionnée
if os.path.exists("/tmp/output"):
    shutil.rmtree("/tmp/output")
result.write.parquet("/tmp/output")

baseline_time = time.time() - start
print(f"\n⏱️ Temps baseline : {baseline_time:.2f}s")

# ✅ VERSION OPTIMISÉE
print("="*50)
print("✅ PIPELINE OPTIMISÉ")
print("="*50)

# Activer AQE
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

start = time.time()

# 1. Lecture avec schema explicite pour CSV
from pyspark.sql.types import StructType, StructField, IntegerType, StringType, DoubleType

products_schema = StructType([
    StructField("product_id", IntegerType()),
    StructField("product_name", StringType()),
    StructField("category", StringType()),
    StructField("base_price", DoubleType())
])

transactions_df = spark.read.parquet("/tmp/transactions")
products_df = spark.read.option("header", True).schema(products_schema).csv("/tmp/products")  # ✅ Schema explicite
customers_df = spark.read.parquet("/tmp/customers")

# 2. Broadcast les petites tables
products_df = broadcast(products_df)

# 3. Expression native au lieu de UDF
amount_category_expr = when(col("amount") > 500, "high") \
    .when(col("amount") > 100, "medium") \
    .otherwise("low")

# 4. Pipeline optimisé
result_optimized = transactions_df \
    .join(products_df, "product_id") \
    .join(customers_df, "customer_id") \
    .withColumn("amount_category", amount_category_expr) \
    .withColumn("year", year("date")) \
    .withColumn("month", month("date")) \
    .groupBy("category", "segment", "amount_category", "year", "month") \
    .agg(
        count("*").alias("num_transactions"),
        sum("amount").alias("total_amount")
    )

# 5. Écriture partitionnée
if os.path.exists("/tmp/output_optimized"):
    shutil.rmtree("/tmp/output_optimized")
result_optimized.write.partitionBy("year", "month").parquet("/tmp/output_optimized")

optimized_time = time.time() - start
print(f"\n⏱️ Temps optimisé : {optimized_time:.2f}s")

# Résumé des optimisations
print("\n" + "="*60)
print("📊 RÉSUMÉ DES OPTIMISATIONS")
print("="*60)

speedup = baseline_time / optimized_time if optimized_time > 0 else 0
reduction = (1 - optimized_time / baseline_time) * 100 if baseline_time > 0 else 0

print(f"""
⏱️ Temps baseline    : {baseline_time:.2f}s
⏱️ Temps optimisé    : {optimized_time:.2f}s
🚀 Speedup           : {speedup:.1f}x
📉 Réduction         : {reduction:.0f}%

Optimisations appliquées :
  ✅ AQE activé (Adaptive Query Execution)
  ✅ Schema explicite pour CSV (pas d'inférence)
  ✅ Broadcast join pour products (10K lignes)
  ✅ Expression native au lieu de UDF Python
  ✅ Écriture partitionnée par year/month
""")

# Nettoyage
spark.stop()
print("✅ SparkSession arrêtée")

# Nettoyage des fichiers temporaires (optionnel)
# import shutil
# for path in ["/tmp/transactions", "/tmp/products", "/tmp/customers", 
#              "/tmp/output", "/tmp/output_optimized", "/tmp/test_parquet", "/tmp/partitioned_data"]:
#     if os.path.exists(path):
#         shutil.rmtree(path)
# print("🧹 Fichiers temporaires supprimés")