from pathlib import Path
from shutil import rmtree

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-parquet-read-write")
    .master("local[*]")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

base_path = Path("spark-parquet-demo")
source_path = base_path / "input" / "orders"
output_path = base_path / "output" / "paid-orders"

if base_path.exists():
    rmtree(base_path)

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

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

(
    seed_orders
    .write
    .mode("overwrite")
    .partitionBy("region")
    .parquet(str(source_path))
)

orders = spark.read.parquet(str(source_path))
paid_orders = (
    orders
    .where(col("status") == "paid")
    .select("order_id", "region", "amount", "order_date")
)

print("Source row count:", orders.count())
orders.printSchema()
paid_orders.orderBy("order_id").show(truncate=False)

(
    paid_orders
    .coalesce(1)
    .write
    .mode("overwrite")
    .option("compression", "snappy")
    .partitionBy("region")
    .parquet(str(output_path))
)

read_back = spark.read.parquet(str(output_path)).orderBy("order_id")

print("Read-back row count:", read_back.count())
read_back.printSchema()
read_back.show(truncate=False)

spark.stop()