from pyspark.sql import SparkSession, functions as F

spark = (
    SparkSession.builder
    .appName("sg-shuffle-partitions-check")
    .master("local[2]")
    .config("spark.ui.enabled", "false")
    .config("spark.ui.showConsoleProgress", "false")
    .config("spark.sql.adaptive.enabled", "false")
    .config("spark.sql.shuffle.partitions", "6")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

orders = spark.range(0, 60, 1, numPartitions=3).select(
    (F.col("id") % 6).alias("bucket")
)

totals = orders.groupBy("bucket").count().orderBy("bucket")
rows = totals.collect()

print(f"spark.sql.shuffle.partitions = {spark.conf.get('spark.sql.shuffle.partitions')}")
print(f"result partitions = {totals.rdd.getNumPartitions()}")
print("Physical plan")
totals.explain(mode="simple")
print("rows = " + ", ".join(f"{row.bucket}:{row['count']}" for row in rows))

spark.stop()