Video Script: Spark Pools Deep Dive¶

# Linked to Azure Key Vault
Environment Variables:
  STORAGE_ACCOUNT_KEY: @Microsoft.KeyVault(SecretUri=https://myvault.vault.azure.net/secrets/storagekey/)
  API_ENDPOINT: @Microsoft.KeyVault(SecretUri=https://myvault.vault.azure.net/secrets/apiurl/)

# Method 1: Using ABFSS protocol with account key
spark.conf.set(
    "fs.azure.account.key.mydatalake.dfs.core.windows.net",
    "your-account-key"  # Use Key Vault in production!
)

# Read Parquet files
df_parquet = spark.read.parquet(
    "abfss://container@mydatalake.dfs.core.windows.net/data/sales/*.parquet"
)

# Method 2: Using managed identity (recommended)
# No credentials needed - uses pool's managed identity
df_secure = spark.read.parquet(
    "abfss://container@mydatalake.dfs.core.windows.net/data/customers/"
)

# Method 3: Read CSV with schema inference
df_csv = spark.read \
    .option("header", "true") \
    .option("inferSchema", "true") \
    .csv("abfss://container@mydatalake.dfs.core.windows.net/data/products.csv")

# Display schema and preview
df_parquet.printSchema()
display(df_parquet.limit(10))

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# Filter and aggregate
sales_summary = df_parquet \
    .filter(F.col("sale_date") >= "2024-01-01") \
    .groupBy("product_category", "region") \
    .agg(
        F.sum("sales_amount").alias("total_sales"),
        F.count("*").alias("transaction_count"),
        F.avg("sales_amount").alias("avg_transaction_value"),
        F.max("sales_amount").alias("max_transaction"),
        F.min("sales_amount").alias("min_transaction")
    )

# Window function for running totals
window_spec = Window.partitionBy("product_category").orderBy("region")

enriched_summary = sales_summary.withColumn(
    "running_total",
    F.sum("total_sales").over(window_spec)
).withColumn(
    "pct_of_category",
    F.col("total_sales") / F.sum("total_sales").over(Window.partitionBy("product_category")) * 100
)

display(enriched_summary)

# Join with customer dimension
df_customers = spark.read.parquet(
    "abfss://container@mydatalake.dfs.core.windows.net/dimensions/customers/"
)

# Broadcast join for small dimension tables (performance optimization)
from pyspark.sql.functions import broadcast

enriched_sales = df_parquet \
    .join(
        broadcast(df_customers),
        df_parquet.customer_id == df_customers.customer_id,
        "left"
    ) \
    .select(
        df_parquet["*"],
        df_customers.customer_name,
        df_customers.customer_tier,
        df_customers.lifetime_value
    )

# Verify join results
print(f"Original records: {df_parquet.count():,}")
print(f"After join: {enriched_sales.count():,}")

display(enriched_sales.limit(20))

# Write data as Delta Lake table
delta_path = "abfss://container@mydatalake.dfs.core.windows.net/delta/sales_enriched"

enriched_sales.write \
    .format("delta") \
    .mode("overwrite") \
    .option("overwriteSchema", "true") \
    .partitionBy("sale_date") \
    .save(delta_path)

print("Data written to Delta Lake successfully!")

# Create table in metastore for SQL access
enriched_sales.write \
    .format("delta") \
    .mode("overwrite") \
    .option("overwriteSchema", "true") \
    .partitionBy("sale_date") \
    .saveAsTable("default.sales_enriched")

from delta.tables import DeltaTable

# Load Delta table
delta_table = DeltaTable.forPath(spark, delta_path)

# View history
history_df = delta_table.history()
display(history_df.select("version", "timestamp", "operation", "operationMetrics"))

# Time travel - query previous versions
df_v0 = spark.read.format("delta").option("versionAsOf", 0).load(delta_path)
df_yesterday = spark.read.format("delta") \
    .option("timestampAsOf", "2024-01-15 00:00:00") \
    .load(delta_path)

print(f"Current version count: {spark.read.format('delta').load(delta_path).count()}")
print(f"Version 0 count: {df_v0.count()}")

from delta.tables import DeltaTable

# Create updates DataFrame
updates_df = spark.createDataFrame([
    (1001, "Product A", 150.00, "2024-01-20"),
    (1002, "Product B", 250.00, "2024-01-20"),
    (1003, "Product C", 350.00, "2024-01-20"),
], ["product_id", "product_name", "price", "update_date"])

# Load target Delta table
target_table = DeltaTable.forPath(spark, delta_path)

# Perform merge (upsert)
target_table.alias("target").merge(
    updates_df.alias("updates"),
    "target.product_id = updates.product_id"
).whenMatchedUpdate(set={
    "product_name": "updates.product_name",
    "price": "updates.price",
    "update_date": "updates.update_date"
}).whenNotMatchedInsert(values={
    "product_id": "updates.product_id",
    "product_name": "updates.product_name",
    "price": "updates.price",
    "update_date": "updates.update_date"
}).execute()

print("Merge operation completed!")

# Check merge metrics
display(target_table.history(1).select("operationMetrics"))

# View physical execution plan
enriched_sales.explain(mode="extended")

# Or use formatted explain
enriched_sales.explain(mode="formatted")

# Cost-based optimization
spark.conf.set("spark.sql.cbo.enabled", "true")
spark.sql("ANALYZE TABLE sales_enriched COMPUTE STATISTICS FOR ALL COLUMNS")

# Check current partitions
print(f"Current partitions: {df_parquet.rdd.getNumPartitions()}")

# Repartition for better parallelism
df_repartitioned = df_parquet.repartition(40, "product_category")

print(f"After repartition: {df_repartitioned.rdd.getNumPartitions()}")

# Coalesce to reduce partitions (no shuffle)
df_coalesced = df_repartitioned.coalesce(10)

print(f"After coalesce: {df_coalesced.rdd.getNumPartitions()}")

# Partition size check
spark.sql("""
    SELECT
        input_file_name(),
        COUNT(*) as record_count,
        SUM(sales_amount) as total_amount
    FROM sales_enriched
    GROUP BY input_file_name()
    ORDER BY record_count DESC
""").show(truncate=False)

# Rule of thumb: 2-3 partitions per CPU core
num_executors = 10
cores_per_executor = 4
optimal_partitions = num_executors * cores_per_executor * 2

print(f"Recommended partitions: {optimal_partitions}")

# Adaptive Query Execution handles this automatically in Spark 3.x
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

from pyspark.storagelevel import StorageLevel

# Cache in memory only
df_parquet.cache()

# Or specify storage level
df_parquet.persist(StorageLevel.MEMORY_AND_DISK)

# Perform action to materialize cache
df_parquet.count()

# Check cache usage
spark.catalog.cacheTable("sales_enriched")

# View cached tables
for table in spark.catalog.listTables():
    if table.isCached:
        print(f"Cached table: {table.name}")

# Clear cache when done
df_parquet.unpersist()
spark.catalog.uncacheTable("sales_enriched")

# Get current Spark configuration
configs = spark.sparkContext.getConf().getAll()
for conf in sorted(configs):
    if 'executor' in conf[0] or 'driver' in conf[0]:
        print(f"{conf[0]}: {conf[1]}")

# Monitor executor metrics
sc = spark.sparkContext
print(f"Active executors: {len(sc._jsc.sc().statusTracker().getExecutorInfos())}")
print(f"Default parallelism: {sc.defaultParallelism}")

from pyspark.ml.feature import VectorAssembler, StringIndexer, StandardScaler
from pyspark.ml import Pipeline

# Prepare features
feature_cols = ["unit_price", "quantity", "discount_pct", "customer_tier_index"]

# Index categorical variables
indexer = StringIndexer(inputCol="customer_tier", outputCol="customer_tier_index")

# Assemble features
assembler = VectorAssembler(
    inputCols=feature_cols,
    outputCol="features_raw"
)

# Scale features
scaler = StandardScaler(
    inputCol="features_raw",
    outputCol="features",
    withStd=True,
    withMean=True
)

# Create pipeline
pipeline = Pipeline(stages=[indexer, assembler, scaler])

# Fit and transform
model = pipeline.fit(enriched_sales)
df_prepared = model.transform(enriched_sales)

display(df_prepared.select("features", "total_amount").limit(10))

from pyspark.ml.regression import LinearRegression
from pyspark.ml.evaluation import RegressionEvaluator

# Split data
train_df, test_df = df_prepared.randomSplit([0.8, 0.2], seed=42)

print(f"Training set: {train_df.count():,} records")
print(f"Test set: {test_df.count():,} records")

# Train linear regression model
lr = LinearRegression(
    featuresCol="features",
    labelCol="total_amount",
    maxIter=100,
    regParam=0.3,
    elasticNetParam=0.8
)

lr_model = lr.fit(train_df)

# Model summary
print(f"Coefficients: {lr_model.coefficients}")
print(f"Intercept: {lr_model.intercept}")
print(f"RMSE: {lr_model.summary.rootMeanSquaredError}")
print(f"R2: {lr_model.summary.r2}")

# Make predictions
predictions = lr_model.transform(test_df)

# Evaluate model
evaluator = RegressionEvaluator(
    labelCol="total_amount",
    predictionCol="prediction",
    metricName="rmse"
)

rmse = evaluator.evaluate(predictions)
print(f"Test RMSE: {rmse:,.2f}")

# Show predictions
display(predictions.select("total_amount", "prediction", "features").limit(20))

# Save model
model_path = "abfss://container@mydatalake.dfs.core.windows.net/models/sales_prediction"
lr_model.write().overwrite().save(model_path)

print("Model saved successfully!")

# Symptoms:
# java.lang.OutOfMemoryError: Java heap space
# org.apache.spark.shuffle.FetchFailedException

# Solutions:
# 1. Increase executor memory
spark.conf.set("spark.executor.memory", "16g")
spark.conf.set("spark.executor.memoryOverhead", "4g")

# 2. Increase partitions
df_large = df_large.repartition(200)

# 3. Use iterative processing
from pyspark.sql.functions import col

def process_in_batches(df, batch_size=1000000):
    total_count = df.count()
    for offset in range(0, total_count, batch_size):
        batch = df.limit(batch_size).offset(offset)
        # Process batch
        yield batch

# Enable adaptive query execution
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")

# Use broadcast joins for small tables
from pyspark.sql.functions import broadcast
result = large_df.join(broadcast(small_df), "key")

# Enable columnar storage
spark.conf.set("spark.sql.parquet.enableVectorizedReader", "true")

# Detect skew
df.groupBy("partition_key").count().orderBy(F.desc("count")).show()

# Solution 1: Salt the key
from pyspark.sql.functions import rand, concat

df_salted = df.withColumn("salt", (rand() * 10).cast("int"))
df_salted = df_salted.withColumn("salted_key", concat("partition_key", "salt"))

# Solution 2: Enable skew join optimization
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionFactor", "5")

# Performance
✅ Use Parquet or Delta for storage
✅ Partition large tables by date
✅ Enable adaptive query execution
✅ Cache only reused DataFrames
✅ Use broadcast for small tables
✅ Monitor Spark UI regularly

# Cost Optimization
✅ Enable auto-pause (15 minutes)
✅ Right-size executor nodes
✅ Use autoscaling
✅ Clean up unused data
✅ Use Serverless SQL for ad-hoc queries

# Code Quality
✅ Use consistent naming conventions
✅ Add comments for complex logic
✅ Implement error handling
✅ Log important operations
✅ Version control notebooks