from pyspark.sql import SparkSession
from pyspark.sql.functions import col

spark = (
    SparkSession.builder.appName("sg-aqe-check")
    .config("spark.sql.adaptive.enabled", "true")
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
    .config("spark.sql.shuffle.partitions", "8")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

orders = spark.range(0, 1000).select(
    (col("id") % 8).alias("customer_id"),
    (col("id") * 10).alias("amount"),
)
customers = spark.range(0, 8).select(col("id").alias("customer_id"))

result = (
    orders.join(customers, "customer_id")
    .groupBy("customer_id")
    .count()
    .orderBy("customer_id")
)

rows = result.collect()

print(f"spark.sql.adaptive.enabled = {spark.conf.get('spark.sql.adaptive.enabled')}")
print(
    "spark.sql.adaptive.coalescePartitions.enabled = "
    f"{spark.conf.get('spark.sql.adaptive.coalescePartitions.enabled')}"
)
print(f"spark.sql.shuffle.partitions = {spark.conf.get('spark.sql.shuffle.partitions')}")
result.explain()
print(f"row_count = {len(rows)}")
print("first_rows = " + ", ".join(f"{row.customer_id}:{row['count']}" for row in rows[:3]))

spark.stop()