from pyspark.sql import SparkSession
from pyspark.sql import functions as F


spark = (
    SparkSession.builder.master("local[1]")
    .appName("spark-csv-read-write")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

schema = (
    "order_id STRING, region STRING, "
    "item_count INT, order_total DOUBLE, notes STRING"
)

orders = (
    spark.read.schema(schema)
    .option("header", True)
    .option("enforceSchema", False)
    .option("mode", "FAILFAST")
    .csv("orders.csv")
)

priority_orders = (
    orders.where(
        (F.col("region") == "APAC")
        & (F.col("item_count") >= 3)
    )
    .select(
        "order_id",
        "region",
        "item_count",
        "order_total",
        "notes",
    )
    .orderBy("order_id")
)

output_path = "output/csv-priority-orders"

(
    priority_orders.coalesce(1)
    .write.mode("overwrite")
    .option("header", True)
    .csv(output_path)
)

read_back = (
    spark.read.schema(priority_orders.schema)
    .option("header", True)
    .option("enforceSchema", False)
    .csv(output_path)
    .orderBy("order_id")
)

assert read_back.collect() == priority_orders.collect()

print("Input rows:", orders.count())
print("Read-back rows:", read_back.count())
read_back.show(truncate=False)

spark.stop()