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

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

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"),
], schema)

paid_orders = (
    orders
    .where(col("status") == "paid")
    .select("order_id", "region", "amount", "order_date")
)

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

(
    paid_orders
    .coalesce(1)
    .write
    .mode("overwrite")
    .option("compression", "snappy")
    .partitionBy("region")
    .orc("output/orc-sales")
)

read_back = spark.read.orc("output/orc-sales").orderBy("order_id")

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

spark.stop()