import shutil

from pyspark.sql import SparkSession

base_path = "/tmp/sg-spark-package"
avro_path = f"{base_path}/orders-avro"
shutil.rmtree(base_path, ignore_errors=True)

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

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

orders.write.format("avro").mode("overwrite").save(avro_path)
loaded = spark.read.format("avro").load(avro_path)

print(f"avro_rows={loaded.count()}")
loaded.orderBy("order_id").show(truncate=False)

spark.stop()
shutil.rmtree(base_path, ignore_errors=True)