Spark SQL Deep Dive & Optimisation

from pyspark.sql import SparkSession
from pyspark.sql.functions import *
from pyspark.sql.window import Window
from pyspark.sql.types import *

# Créer une SparkSession
spark = SparkSession.builder \
    .appName("Spark SQL Deep Dive") \
    .config("spark.sql.adaptive.enabled", "true") \
    .getOrCreate()

print(f"✅ Spark {spark.version} initialisé")
print(f"🔍 Spark UI : {spark.sparkContext.uiWebUrl}")

# Créer des données de test
sales_data = [
    (1, "Electronics", "North", 1200, "2024-01-15"),
    (2, "Electronics", "South", 800, "2024-01-16"),
    (3, "Clothing", "North", 450, "2024-01-15"),
    (4, "Clothing", "South", 650, "2024-01-17"),
    (5, "Electronics", "North", 950, "2024-01-18"),
    (6, "Food", "South", 320, "2024-01-15"),
]

sales_df = spark.createDataFrame(sales_data, 
    ["id", "category", "region", "amount", "date"])

# Créer une vue temporaire
sales_df.createOrReplaceTempView("sales")

# Utiliser SQL
result = spark.sql("""
    SELECT category, SUM(amount) as total_sales
    FROM sales
    WHERE amount > 500
    GROUP BY category
    ORDER BY total_sales DESC
""")

result.show()

# Équivalence DataFrame API vs SQL

# === DataFrame API ===
result_df = sales_df \
    .filter(col("amount") > 500) \
    .groupBy("category") \
    .agg(sum("amount").alias("total_sales")) \
    .orderBy(desc("total_sales"))

# === SQL ===
result_sql = spark.sql("""
    SELECT category, SUM(amount) as total_sales
    FROM sales
    WHERE amount > 500
    GROUP BY category
    ORDER BY total_sales DESC
""")

print("DataFrame API:")
result_df.show()

print("SQL:")
result_sql.show()

# Les plans sont identiques !
print("\n✅ Les deux approches génèrent le même plan d'exécution")

# Créer des données pour démontrer les joins
products_data = [
    (1, "Laptop", "Electronics"),
    (2, "T-Shirt", "Clothing"),
    (3, "Apple", "Food"),
]
products_df = spark.createDataFrame(products_data, ["product_id", "name", "category"])
products_df.createOrReplaceTempView("products")

# Voir le plan d'exécution
spark.sql("""
    EXPLAIN FORMATTED
    SELECT s.*, p.name
    FROM sales s
    JOIN products p ON s.category = p.category
    WHERE s.amount > 500
""").show(truncate=False)

# Créer des données plus riches pour les window functions
orders_data = [
    (1, 101, "2024-01-01", 150.0, "Premium"),
    (2, 101, "2024-01-15", 200.0, "Premium"),
    (3, 101, "2024-02-01", 180.0, "Premium"),
    (4, 102, "2024-01-05", 300.0, "Standard"),
    (5, 102, "2024-01-20", 250.0, "Standard"),
    (6, 103, "2024-01-10", 400.0, "Premium"),
    (7, 103, "2024-01-25", 400.0, "Premium"),  # Ex-aequo intentionnel
    (8, 103, "2024-02-10", 350.0, "Premium"),
    (9, 104, "2024-01-03", 100.0, "Standard"),
    (10, 104, "2024-02-15", 120.0, "Standard"),
]

orders_df = spark.createDataFrame(orders_data,
    ["order_id", "customer_id", "order_date", "amount", "segment"])
orders_df = orders_df.withColumn("order_date", to_date(col("order_date")))

orders_df.createOrReplaceTempView("orders")
orders_df.show()

# Démonstration des fonctions de ranking
ranking_result = spark.sql("""
    SELECT 
        customer_id,
        order_date,
        amount,
        ROW_NUMBER() OVER (PARTITION BY customer_id ORDER BY amount DESC) as row_num,
        RANK() OVER (PARTITION BY customer_id ORDER BY amount DESC) as rank,
        DENSE_RANK() OVER (PARTITION BY customer_id ORDER BY amount DESC) as dense_rank,
        NTILE(2) OVER (PARTITION BY customer_id ORDER BY amount DESC) as ntile_2
    FROM orders
    WHERE customer_id = 103
    ORDER BY customer_id, amount DESC
""")

print("📊 Comparaison ROW_NUMBER vs RANK vs DENSE_RANK (customer 103 avec ex-aequo):")
ranking_result.show()

# Cas d'usage : Top 2 commandes par client
top_orders = spark.sql("""
    SELECT * FROM (
        SELECT 
            customer_id,
            order_date,
            amount,
            ROW_NUMBER() OVER (PARTITION BY customer_id ORDER BY amount DESC) as rn
        FROM orders
    )
    WHERE rn <= 2
    ORDER BY customer_id, rn
""")

print("🏆 Top 2 commandes par client:")
top_orders.show()

# Cas d'usage : Délai entre commandes
order_gaps = spark.sql("""
    SELECT 
        customer_id,
        order_date,
        amount,
        LAG(order_date, 1) OVER (PARTITION BY customer_id ORDER BY order_date) as prev_order_date,
        LEAD(order_date, 1) OVER (PARTITION BY customer_id ORDER BY order_date) as next_order_date,
        DATEDIFF(
            order_date, 
            LAG(order_date, 1) OVER (PARTITION BY customer_id ORDER BY order_date)
        ) as days_since_last_order
    FROM orders
    ORDER BY customer_id, order_date
""")

print("📅 Analyse du délai entre commandes:")
order_gaps.show()

# FIRST_VALUE et LAST_VALUE
first_last = spark.sql("""
    SELECT 
        customer_id,
        order_date,
        amount,
        FIRST_VALUE(order_date) OVER (
            PARTITION BY customer_id ORDER BY order_date
        ) as first_order_date,
        FIRST_VALUE(amount) OVER (
            PARTITION BY customer_id ORDER BY order_date
        ) as first_order_amount
    FROM orders
    ORDER BY customer_id, order_date
""")

print("📊 FIRST_VALUE - Date et montant de la première commande:")
first_last.show()

# Running total, Running average, % du total
window_agg = spark.sql("""
    SELECT 
        customer_id,
        order_date,
        amount,
        
        -- Somme cumulative
        SUM(amount) OVER (
            PARTITION BY customer_id 
            ORDER BY order_date
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        ) as running_total,
        
        -- Total du client (pour calculer %)
        SUM(amount) OVER (PARTITION BY customer_id) as customer_total,
        
        -- % de chaque commande
        ROUND(amount / SUM(amount) OVER (PARTITION BY customer_id) * 100, 1) as pct_of_customer_total
        
    FROM orders
    ORDER BY customer_id, order_date
""")

print("📊 Agrégations fenêtrées:")
window_agg.show()

# Rolling average (moyenne mobile)
# Créer plus de données pour démontrer
daily_sales = spark.createDataFrame([
    ("2024-01-01", 100), ("2024-01-02", 150), ("2024-01-03", 120),
    ("2024-01-04", 180), ("2024-01-05", 90), ("2024-01-06", 200),
    ("2024-01-07", 170), ("2024-01-08", 140), ("2024-01-09", 160),
    ("2024-01-10", 190),
], ["date", "sales"])
daily_sales = daily_sales.withColumn("date", to_date(col("date")))
daily_sales.createOrReplaceTempView("daily_sales")

# Moyenne mobile sur 3 jours
rolling = spark.sql("""
    SELECT 
        date,
        sales,
        ROUND(AVG(sales) OVER (
            ORDER BY date
            ROWS BETWEEN 2 PRECEDING AND CURRENT ROW
        ), 1) as rolling_avg_3d,
        SUM(sales) OVER (
            ORDER BY date
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        ) as cumulative_sum
    FROM daily_sales
""")

print("📈 Moyenne mobile 3 jours et somme cumulative:")
rolling.show()

# Données mensuelles
monthly_data = spark.createDataFrame([
    ("Alice", "Jan", 1000), ("Alice", "Feb", 1200), ("Alice", "Mar", 1100),
    ("Bob", "Jan", 800), ("Bob", "Feb", 900), ("Bob", "Mar", 950),
    ("Charlie", "Jan", 1500), ("Charlie", "Feb", 1400), ("Charlie", "Mar", 1600),
], ["salesperson", "month", "revenue"])

monthly_data.createOrReplaceTempView("monthly_sales")

print("Données originales (format long):")
monthly_data.show()

# PIVOT : transformer les mois en colonnes
pivoted = spark.sql("""
    SELECT * FROM monthly_sales
    PIVOT (
        SUM(revenue)
        FOR month IN ('Jan', 'Feb', 'Mar')
    )
""")

print("Après PIVOT (format large):")
pivoted.show()

# PIVOT avec DataFrame API
pivoted_df = monthly_data.groupBy("salesperson").pivot("month", ["Jan", "Feb", "Mar"]).sum("revenue")

print("PIVOT avec DataFrame API:")
pivoted_df.show()

# Créer une vue du DataFrame pivoté
pivoted_df.createOrReplaceTempView("pivoted_sales")

# UNPIVOT (Spark 3.4+)
# Pour les versions antérieures, utiliser stack()
unpivoted = spark.sql("""
    SELECT 
        salesperson,
        stack(3, 
            'Jan', Jan, 
            'Feb', Feb, 
            'Mar', Mar
        ) as (month, revenue)
    FROM pivoted_sales
""")

print("Après UNPIVOT (retour au format long):")
unpivoted.show()

# GROUPING SETS : plusieurs agrégations en une requête
grouping_result = spark.sql("""
    SELECT 
        category,
        region,
        SUM(amount) as total_sales,
        COUNT(*) as num_transactions
    FROM sales
    GROUP BY GROUPING SETS (
        (category, region),  -- Par catégorie ET région
        (category),          -- Par catégorie seulement
        (region),            -- Par région seulement
        ()                   -- Total global
    )
    ORDER BY category NULLS LAST, region NULLS LAST
""")

print("📊 GROUPING SETS - Agrégations multiples:")
grouping_result.show()

# ROLLUP : hiérarchie (catégorie → région → total)
rollup_result = spark.sql("""
    SELECT 
        category,
        region,
        SUM(amount) as total_sales
    FROM sales
    GROUP BY ROLLUP (category, region)
    ORDER BY category NULLS LAST, region NULLS LAST
""")

print("📊 ROLLUP - Hiérarchie d'agrégations:")
rollup_result.show()

# CUBE : toutes les combinaisons possibles
cube_result = spark.sql("""
    SELECT 
        category,
        region,
        SUM(amount) as total_sales
    FROM sales
    GROUP BY CUBE (category, region)
    ORDER BY category NULLS LAST, region NULLS LAST
""")

print("📊 CUBE - Toutes les combinaisons:")
cube_result.show()

# GROUPING() : identifier les sous-totaux
grouping_with_flags = spark.sql("""
    SELECT 
        category,
        region,
        SUM(amount) as total_sales,
        GROUPING(category) as is_category_subtotal,
        GROUPING(region) as is_region_subtotal,
        CASE 
            WHEN GROUPING(category) = 1 AND GROUPING(region) = 1 THEN 'GRAND TOTAL'
            WHEN GROUPING(category) = 1 THEN 'Region Subtotal'
            WHEN GROUPING(region) = 1 THEN 'Category Subtotal'
            ELSE 'Detail'
        END as row_type
    FROM sales
    GROUP BY CUBE (category, region)
    ORDER BY GROUPING(category), GROUPING(region), category, region
""")

print("📊 GROUPING() pour identifier les sous-totaux:")
grouping_with_flags.show()

# Créer des données avec arrays
customers_with_tags = spark.createDataFrame([
    (1, "Alice", ["premium", "loyal", "newsletter"]),
    (2, "Bob", ["new", "newsletter"]),
    (3, "Charlie", ["premium", "vip"]),
], ["id", "name", "tags"])

customers_with_tags.createOrReplaceTempView("customers_tags")

print("Données avec arrays:")
customers_with_tags.show(truncate=False)

# EXPLODE : une ligne par tag
exploded = spark.sql("""
    SELECT id, name, tag
    FROM customers_tags
    LATERAL VIEW EXPLODE(tags) t AS tag
""")

print("Après EXPLODE:")
exploded.show()

# POSEXPLODE : avec position
posexploded = spark.sql("""
    SELECT id, name, pos, tag
    FROM customers_tags
    LATERAL VIEW POSEXPLODE(tags) t AS pos, tag
""")

print("POSEXPLODE (avec index):")
posexploded.show()

# EXPLODE avec Map
customers_with_attrs = spark.createDataFrame([
    (1, "Alice", {"city": "Paris", "country": "France"}),
    (2, "Bob", {"city": "London", "country": "UK", "postal": "SW1"}),
], ["id", "name", "attributes"])

customers_with_attrs.createOrReplaceTempView("customers_attrs")

print("Données avec Map:")
customers_with_attrs.show(truncate=False)

# EXPLODE sur Map → (key, value)
exploded_map = spark.sql("""
    SELECT id, name, key, value
    FROM customers_attrs
    LATERAL VIEW EXPLODE(attributes) t AS key, value
""")

print("EXPLODE sur Map:")
exploded_map.show()

# Données JSON
events = spark.createDataFrame([
    (1, '{"user": "alice", "action": "click", "details": {"page": "home", "duration": 5}}'),
    (2, '{"user": "bob", "action": "purchase", "details": {"page": "cart", "amount": 99.99}}'),
], ["id", "json_data"])

events.createOrReplaceTempView("events")

# Extraire des champs JSON
json_extracted = spark.sql("""
    SELECT 
        id,
        get_json_object(json_data, '$.user') as user,
        get_json_object(json_data, '$.action') as action,
        get_json_object(json_data, '$.details.page') as page
    FROM events
""")

print("Extraction JSON avec get_json_object:")
json_extracted.show()

# Parser JSON complet avec schema
from pyspark.sql.types import StructType, StructField, StringType, DoubleType

details_schema = StructType([
    StructField("page", StringType()),
    StructField("duration", DoubleType()),
    StructField("amount", DoubleType())
])

event_schema = StructType([
    StructField("user", StringType()),
    StructField("action", StringType()),
    StructField("details", details_schema)
])

parsed = events.withColumn("parsed", from_json(col("json_data"), event_schema))
parsed.select("id", "parsed.user", "parsed.action", "parsed.details.page").show()

# CTE : structurer une requête complexe
cte_result = spark.sql("""
    WITH 
    -- Étape 1 : Agrégation par client
    customer_stats AS (
        SELECT 
            customer_id,
            COUNT(*) as num_orders,
            SUM(amount) as total_spent,
            AVG(amount) as avg_order
        FROM orders
        GROUP BY customer_id
    ),
    
    -- Étape 2 : Moyenne globale
    global_avg AS (
        SELECT AVG(total_spent) as avg_total_spent
        FROM customer_stats
    )
    
    -- Requête finale : clients au-dessus de la moyenne
    SELECT 
        c.*,
        g.avg_total_spent,
        CASE WHEN c.total_spent > g.avg_total_spent THEN 'Above Average' ELSE 'Below Average' END as status
    FROM customer_stats c
    CROSS JOIN global_avg g
    ORDER BY total_spent DESC
""")

print("📊 Analyse client avec CTEs:")
cte_result.show()

# ❌ Subquery corrélée (potentiellement lent)
# SELECT * FROM orders o
# WHERE amount > (SELECT AVG(amount) FROM orders WHERE customer_id = o.customer_id)

# ✅ Window function (plus efficace)
efficient_query = spark.sql("""
    SELECT * FROM (
        SELECT 
            *,
            AVG(amount) OVER (PARTITION BY customer_id) as avg_customer_amount
        FROM orders
    )
    WHERE amount > avg_customer_amount
""")

print("✅ Commandes au-dessus de la moyenne du client (window function):")
efficient_query.show()

# Créer des tables pour démontrer les hints
spark.sql("DROP TABLE IF EXISTS dim_segment")
segments = spark.createDataFrame([
    ("Premium", 1.2), ("Standard", 1.0)
], ["segment", "multiplier"])
segments.createOrReplaceTempView("dim_segment")

# BROADCAST hint
broadcast_join = spark.sql("""
    SELECT /*+ BROADCAST(dim_segment) */ 
        o.*,
        s.multiplier,
        o.amount * s.multiplier as adjusted_amount
    FROM orders o
    JOIN dim_segment s ON o.segment = s.segment
""")

print("Plan avec BROADCAST hint:")
broadcast_join.explain()
print("\nRésultat:")
broadcast_join.show(5)

# Collecter les statistiques (améliore l'optimiseur)
# Note : fonctionne sur des tables persistées, pas des vues temporaires

print("💡 Pour collecter des statistiques sur une table persistée :")
print("""
-- Statistiques de base
ANALYZE TABLE my_table COMPUTE STATISTICS

-- Statistiques sur colonnes spécifiques
ANALYZE TABLE my_table COMPUTE STATISTICS FOR COLUMNS col1, col2

-- Voir les statistiques
DESCRIBE EXTENDED my_table
""")

# Approximation vs COUNT DISTINCT
comparison = spark.sql("""
    SELECT 
        COUNT(DISTINCT customer_id) as exact_count,
        APPROX_COUNT_DISTINCT(customer_id) as approx_count
    FROM orders
""")

print("COUNT DISTINCT vs APPROX_COUNT_DISTINCT:")
comparison.show()
print("💡 APPROX_COUNT_DISTINCT est ~10x plus rapide sur de grandes tables")

import os
import shutil

# Créer des données plus riches
orders_extended = spark.createDataFrame([
    # Customer 101 - Premium, 3 commandes
    (1, 101, "2024-01-01", 150.0, "Premium"),
    (2, 101, "2024-01-15", 200.0, "Premium"),
    (3, 101, "2024-02-01", 180.0, "Premium"),
    # Customer 102 - Standard, 4 commandes
    (4, 102, "2024-01-05", 300.0, "Standard"),
    (5, 102, "2024-01-20", 250.0, "Standard"),
    (6, 102, "2024-02-10", 275.0, "Standard"),
    (7, 102, "2024-03-01", 320.0, "Standard"),
    # Customer 103 - Premium, 2 commandes
    (8, 103, "2024-01-10", 400.0, "Premium"),
    (9, 103, "2024-02-15", 450.0, "Premium"),
    # Customer 104 - Standard, 3 commandes
    (10, 104, "2024-01-03", 100.0, "Standard"),
    (11, 104, "2024-01-25", 120.0, "Standard"),
    (12, 104, "2024-02-20", 90.0, "Standard"),
    # Customer 105 - Premium, 1 commande
    (13, 105, "2024-02-01", 500.0, "Premium"),
], ["order_id", "customer_id", "order_date", "amount", "segment"])

orders_extended = orders_extended.withColumn("order_date", to_date(col("order_date")))
orders_extended.createOrReplaceTempView("orders_ext")

print("📦 Données source:")
orders_extended.show()

# Construction du Datamart avec CTEs et Window Functions
datamart = spark.sql("""
    WITH 
    -- Étape 1 : Enrichir les commandes avec métriques temporelles
    orders_enriched AS (
        SELECT 
            customer_id,
            segment,
            order_date,
            amount,
            
            -- Numéro de commande du client
            ROW_NUMBER() OVER (PARTITION BY customer_id ORDER BY order_date) as order_number,
            
            -- Date de la commande précédente
            LAG(order_date) OVER (PARTITION BY customer_id ORDER BY order_date) as prev_order_date,
            
            -- Délai depuis la dernière commande
            DATEDIFF(
                order_date,
                LAG(order_date) OVER (PARTITION BY customer_id ORDER BY order_date)
            ) as days_since_last
        FROM orders_ext
    ),
    
    -- Étape 2 : Agréger par client
    customer_stats AS (
        SELECT 
            customer_id,
            segment,
            
            -- Première commande
            MIN(order_date) as first_order_date,
            
            -- Dernière commande
            MAX(order_date) as last_order_date,
            
            -- Nombre de commandes
            COUNT(*) as total_orders,
            
            -- Revenus
            SUM(amount) as total_revenue,
            AVG(amount) as avg_order_value,
            
            -- Délai moyen entre commandes (exclut NULL de la première commande)
            AVG(days_since_last) as avg_days_between_orders
        FROM orders_enriched
        GROUP BY customer_id, segment
    ),
    
    -- Étape 3 : Ajouter des métriques de segment
    customer_with_segment_stats AS (
        SELECT 
            *,
            
            -- Total du segment
            SUM(total_revenue) OVER (PARTITION BY segment) as segment_total_revenue,
            
            -- % du revenu du segment
            ROUND(total_revenue / SUM(total_revenue) OVER (PARTITION BY segment) * 100, 2) as pct_of_segment,
            
            -- Rang dans le segment
            RANK() OVER (PARTITION BY segment ORDER BY total_revenue DESC) as rank_in_segment,
            
            -- Nombre de clients dans le segment
            COUNT(*) OVER (PARTITION BY segment) as customers_in_segment
        FROM customer_stats
    )
    
    -- Résultat final
    SELECT 
        customer_id,
        segment,
        first_order_date,
        last_order_date,
        total_orders,
        ROUND(total_revenue, 2) as total_revenue,
        ROUND(avg_order_value, 2) as avg_order_value,
        ROUND(avg_days_between_orders, 1) as avg_days_between_orders,
        rank_in_segment,
        pct_of_segment,
        customers_in_segment
    FROM customer_with_segment_stats
    ORDER BY segment, rank_in_segment
""")

print("📊 DATAMART CUSTOMER ANALYTICS:")
datamart.show(truncate=False)

# Exporter le datamart
output_path = "/tmp/customer_datamart"
if os.path.exists(output_path):
    shutil.rmtree(output_path)

datamart.write.partitionBy("segment").parquet(output_path)

print(f"✅ Datamart exporté : {output_path}")
print(f"📁 Partitionné par segment")

# Vérifier la structure
for root, dirs, files in os.walk(output_path):
    level = root.replace(output_path, '').count(os.sep)
    indent = ' ' * 2 * level
    print(f"{indent}{os.path.basename(root)}/")

# Résumé du mini-projet
print("""
╔══════════════════════════════════════════════════════════════╗
║           📊 RÉSUMÉ DU DATAMART                              ║
╠══════════════════════════════════════════════════════════════╣
║                                                              ║
║  Techniques utilisées :                                      ║
║  ✅ CTEs pour structurer la transformation                   ║
║  ✅ ROW_NUMBER() pour numéroter les commandes                ║
║  ✅ LAG() pour calculer le délai entre commandes             ║
║  ✅ RANK() pour classer les clients par segment              ║
║  ✅ SUM() OVER pour calculer les totaux de segment           ║
║  ✅ Partitionnement Parquet par segment                      ║
║                                                              ║
║  Métriques calculées :                                       ║
║  • Date première/dernière commande                           ║
║  • Nombre total de commandes                                 ║
║  • Revenu total et moyen                                     ║
║  • Délai moyen entre commandes                               ║
║  • Rang dans le segment                                      ║
║  • % du revenu du segment                                    ║
║                                                              ║
║  💡 Ce type de transformation sera automatisé                ║
║     avec dbt dans le module 26                               ║
║                                                              ║
╚══════════════════════════════════════════════════════════════╝
""")

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

# Nettoyage des fichiers temporaires (optionnel)
# import shutil
# if os.path.exists("/tmp/customer_datamart"):
#     shutil.rmtree("/tmp/customer_datamart")
# print("🧹 Fichiers temporaires supprimés")