from pathlib import Path
import shutil

from pyspark.sql import SparkSession

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

output = Path("/tmp/sg-spark-output/partition-count")
shutil.rmtree(output, ignore_errors=True)

df = spark.range(0, 16, 1, numPartitions=4)
print(f"original partitions: {df.rdd.getNumPartitions()}")