from pathlib import Path
import shutil

from pyspark.sql import SparkSession
from pyspark.sql.functions import col
from pyspark.sql.types import DoubleType, StringType, StructField, StructType

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

base_path = Path("spark-partitioned-demo")
output = base_path / "output" / "orders"
shutil.rmtree(base_path, ignore_errors=True)

schema = StructType(
    [
        StructField("order_id", StringType(), False),
        StructField("region", StringType(), False),
        StructField("order_date", StringType(), False),
        StructField("status", StringType(), False),
        StructField("amount", DoubleType(), False),
    ]
)

orders = spark.createDataFrame(
    [
        ("ord-1001", "apac", "2026-07-07", "paid", 212.10),
        ("ord-1002", "emea", "2026-07-07", "paid", 149.50),
        ("ord-1003", "emea", "2026-07-08", "cancelled", 42.00),
        ("ord-1004", "na", "2026-07-08", "paid", 87.25),
    ],
    schema,
)

(
    orders
    .coalesce(1)
    .write
    .mode("overwrite")
    .partitionBy("region", "order_date")
    .parquet(str(output))
)

partition_dirs = sorted(
    str(path.relative_to(output))
    for path in output.glob("region=*/order_date=*")
    if path.is_dir()
)

print("Partition directories:")
for directory in partition_dirs:
    print(directory)

read_back = spark.read.parquet(str(output))
print("Read-back row count:", read_back.count())
read_back.orderBy("order_id").show(truncate=False)

apac_paid = read_back.where(
    (col("region") == "apac") & (col("order_date") == "2026-07-07")
)
print("Filtered row count:", apac_paid.count())
apac_paid.select("order_id", "region", "order_date", "amount").show(truncate=False)

spark.stop()