from pathlib import Path

from pyspark.sql import SparkSession


output_path = Path("/tmp/spark-package-check/orders-avro")

spark = (
    SparkSession.builder.master("local[2]")
    .appName("spark-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(str(output_path))
loaded = spark.read.format("avro").load(str(output_path))

expected = [
    ("ORD-1001", "APAC", 3),
    ("ORD-1002", "EMEA", 1),
]
actual = [tuple(row) for row in loaded.orderBy("order_id").collect()]
assert actual == expected, actual

avro_files = list(output_path.glob("part-*.avro"))
assert avro_files, "No Avro data files were written"

print("Package-backed Avro rows:")
loaded.orderBy("order_id").show(truncate=False)
print(f"Avro data files: {len(avro_files)}")

spark.stop()