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

schema = StructType([
    StructField("event_id", StringType(), True),
    StructField("customer", StringType(), True),
    StructField("amount", DoubleType(), True),
    StructField("status", StringType(), True),
    StructField("event_ts", StringType(), True),
])

events = spark.read.schema(schema).json("events.jsonl")
paid_events = (
    events
    .where(col("status") == "paid")
    .select("event_id", "customer", "amount", "event_ts")
)

print("Input row count:", events.count())
events.printSchema()
paid_events.show(truncate=False)

paid_events.coalesce(1).write.mode("overwrite").json("output/json-events")
read_back = spark.read.schema(paid_events.schema).json("output/json-events")

print("Output row count:", read_back.count())
read_back.show(truncate=False)

spark.stop()