import os
import shutil

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

base_path = "/tmp/sg-spark-avro"
input_path = f"{base_path}/orders"
output_path = f"{base_path}/priority-apac"

shutil.rmtree(base_path, ignore_errors=True)

spark = (
    SparkSession.builder.master("local[1]")
    .appName("sg-spark-avro-check")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

orders = spark.createDataFrame(
    [
        ("ORD-1001", "APAC", 3, True),
        ("ORD-1002", "EMEA", 1, False),
        ("ORD-1003", "APAC", 7, True),
    ],
    "order_id string, region string, item_count int, priority boolean",
)

orders.write.format("avro").mode("overwrite").save(input_path)

loaded = spark.read.format("avro").load(input_path)

print("Input Avro schema:")
loaded.printSchema()
print(f"Input rows: {loaded.count()}")

(
    loaded.filter((F.col("region") == "APAC") & F.col("priority"))
    .select("order_id", "region", "item_count")
    .write.format("avro")
    .mode("overwrite")
    .save(output_path)
)

result = spark.read.format("avro").load(output_path)

print("Output Avro rows:")
result.orderBy("order_id").show(truncate=False)
print("Output files:")
for name in sorted(os.listdir(output_path)):
    if not name.startswith("."):
        print(name)

spark.stop()