from pyspark.sql import SparkSession
from pyspark.sql import functions as F


spark = (
    SparkSession.builder.appName("sg-dataframe-explain-plan")
    .master("local[2]")
    .config("spark.ui.enabled", "false")
    .config("spark.sql.adaptive.enabled", "true")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

orders = spark.createDataFrame(
    [
        ("EMEA", "analytics", 1200),
        ("EMEA", "analytics", 800),
        ("APAC", "platform", 500),
        ("EMEA", "platform", 700),
        ("AMER", "analytics", 400),
    ],
    ["region", "team", "amount"],
)

totals = (
    orders.filter(F.col("region") == "EMEA")
    .groupBy("team")
    .agg(F.sum("amount").alias("total_amount"))
    .orderBy(F.col("total_amount").desc())
)

print("Formatted plan")
totals.explain(mode="formatted")

print("Extended plan")
totals.explain(mode="extended")

print("Result rows")
totals.show()

spark.stop()