# Continuous processing — checkpoint every 1 second
query = result.writeStream \
    .format("kafka") \
    .trigger(continuous="1 second") \  # ← continuous mode trigger
    .option("checkpointLocation", "/tmp/ckpt") \
    .start()

# The "1 second" here means: commit progress every 1 second
# NOT that data is processed in 1-second batches
# Records are processed as they arrive (sub-millisecond)

{
  "batchWatermarkMs": 0,
  "batchTimestampMs": 1700000050000,
  "conf": {
    "spark.sql.streaming.stateStore.providerClass": "..."
  }
}
--- v1
{
  "orders": {                       // topic name
    "0": 200,                       // partition 0 → read up to offset 200
    "1": 185,                       // partition 1 → read up to offset 185
    "2": 210                        // partition 2 → read up to offset 210
  }
}

# After query starts, inspect progress
import time
time.sleep(15)  # let a few batches run

progress = query.lastProgress
print(progress)

# Example output structure:
# {
#   "id": "abc-123",
#   "runId": "xyz-456",
#   "batchId": 5,
#   "timestamp": "2024-01-15T10:30:00.000Z",
#   "numInputRows": 1500,            ← how many rows read this batch
#   "inputRowsPerSecond": 150.0,     ← input rate
#   "processedRowsPerSecond": 300.0, ← processing rate
#   "sources": [{
#     "description": "KafkaV2[orders]",
#     "startOffset": {"orders": {"0": 100, "1": 80}},   ← START of batch
#     "endOffset":   {"orders": {"0": 200, "1": 185}},  ← END of batch
#     "numInputRows": 1500,
#     "inputRowsPerSecond": 150.0,
#     "processedRowsPerSecond": 300.0
#   }]
# }

def write_to_delta(batch_df, batch_id):
    # batch_id is automatically provided by foreachBatch
    # We can use it to check if this batch was already written
    
    batch_df.write \
        .format("delta") \
        .mode("append") \
        .option("txnAppId", "my-streaming-app") \
        .option("txnVersion", batch_id) \  # ← Delta uses this for idempotency!
        .save("/delta/orders")

query = result.writeStream \
    .foreachBatch(write_to_delta) \
    .option("checkpointLocation", "/ckpt/orders") \
    .start()

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, from_json
from pyspark.sql.types import StructType, StringType, IntegerType
import time

spark = SparkSession.builder \
    .appName("LifecycleDemo") \
    .config("spark.sql.shuffle.partitions", "4") \
    .getOrCreate()

# Step 1: Define schema
schema = StructType() \
    .add("order_id", StringType()) \
    .add("amount", IntegerType()) \
    .add("status", StringType())

# Step 2: Read from Kafka (source)
kafka_df = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "orders") \
    .option("startingOffsets", "latest") \
    .load()

# Step 3: Apply transformations
orders_df = kafka_df \
    .select(from_json(col("value").cast(StringType()), schema).alias("data")) \
    .select("data.*") \
    .filter(col("status") == "NEW")

# Step 4: Write to Delta sink with checkpoint
query = orders_df.writeStream \
    .format("delta") \
    .outputMode("append") \
    .option("checkpointLocation", "/tmp/ckpt/orders") \  # ← lifecycle tracked here
    .trigger(processingTime="5 seconds") \
    .start("/delta/new_orders")

# Step 5: Monitor lifecycle progress
for i in range(3):
    time.sleep(6)
    p = query.lastProgress
    if p:
        print(f"Batch {p['batchId']}: {p['numInputRows']} rows, "
              f"{p['processedRowsPerSecond']:.1f} rows/sec")

query.awaitTermination()

# The logical plan is compiled ONCE when the query starts
# You can view it with explain() on the streaming DataFrame

df = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "orders") \
    .load()

# View the STREAMING plan (before it runs)
df.select("value").explain(True)

# Output shows:
# == Parsed Logical Plan ==
# Project [value#0]
# +- StreamingRelationV2 ...  ← source node (offsets swapped each batch)
#
# == Analyzed Logical Plan ==  ← compiled once
# == Optimized Logical Plan == ← compiled once
# == Physical Plan ==           ← compiled once, reused per batch

# STATELESS: just filter — no memory needed between batches
stateless_df = orders_df.filter(col("amount") > 100)

# STATEFUL: count per category across ALL batches so far
# Spark must remember previous counts to add new batch's counts
from pyspark.sql.functions import count, sum as spark_sum

stateful_df = orders_df \
    .withWatermark("event_time", "10 minutes") \
    .groupBy("category") \
    .agg(count("*").alias("total_orders"),
         spark_sum("amount").alias("total_revenue"))
# State is: {category -> (count, sum)} persisted across batches

from pyspark.sql import SparkSession
from pyspark.sql.functions import (
    col, from_json, count, sum as spark_sum, to_timestamp
)
from pyspark.sql.types import StructType, StringType, IntegerType

spark = SparkSession.builder \
    .appName("IncrementalPlanDemo") \
    .config("spark.sql.shuffle.partitions", "8") \
    .getOrCreate()

schema = StructType() \
    .add("order_id", StringType()) \
    .add("category", StringType()) \
    .add("amount", IntegerType()) \
    .add("event_time", StringType())

kafka_df = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "orders") \
    .load()

orders_df = kafka_df \
    .select(from_json(col("value").cast(StringType()), schema).alias("d")) \
    .select(
        col("d.category"),
        col("d.amount"),
        to_timestamp(col("d.event_time")).alias("event_time")
    )

# STATEFUL aggregation — plan is compiled once,
# state (category counts) persists across ALL batches
agg_df = orders_df \
    .withWatermark("event_time", "10 minutes") \
    .groupBy("category") \
    .agg(
        count("*").alias("order_count"),
        spark_sum("amount").alias("total_amount")
    )

# Update mode required for aggregations (only changed rows output)
query = agg_df.writeStream \
    .format("delta") \
    .outputMode("update") \
    .option("checkpointLocation", "/tmp/ckpt/category-agg") \
    .trigger(processingTime="30 seconds") \
    .start("/delta/category_summary")

query.awaitTermination()

# List the commit log files
ls -la /tmp/ckpt/orders/commits/
# 0  1  2  3  4    ← batch 4 is the latest committed batch

# View contents of commit file for batch 3
cat /tmp/ckpt/orders/commits/3
# {"nextBatchWatermarkMs":0}
# Simple! Just marks that batch 3 is done + watermark info

# Original query (crashed during batch 5)
query = result.writeStream \
    .format("delta") \
    .option("checkpointLocation", "/tmp/ckpt/orders") \  # ← SAME checkpoint
    .trigger(processingTime="10 seconds") \
    .start("/delta/orders")

# What Spark does on restart:
# 1. Reads commits/  → last committed batch = 4
# 2. Reads offsets/5 → batch 5 was planned with offsets {p0: 200, p1: 185}
# 3. Re-runs batch 5 with EXACT same offset range
# 4. On success, writes commits/5
# 5. Continues with batch 6...

# IMPORTANT: Do NOT change the checkpoint directory between restarts!
# Changing it means starting from scratch (data reprocessed from beginning)

def idempotent_write(batch_df, batch_id):
    """
    This function will be called AGAIN if batch_id is replayed.
    Using MERGE ensures no duplicates.
    """
    from delta.tables import DeltaTable
    
    # Create target table if not exists
    target_path = "/delta/orders_processed"
    
    if DeltaTable.isDeltaTable(spark, target_path):
        delta_tbl = DeltaTable.forPath(spark, target_path)
        # MERGE deduplicates — if order_id exists, update; else insert
        delta_tbl.alias("target") \
            .merge(batch_df.alias("source"), "target.order_id = source.order_id") \
            .whenMatchedUpdateAll() \
            .whenNotMatchedInsertAll() \
            .execute()
    else:
        # First time — just write
        batch_df.write.format("delta").save(target_path)

query = orders_df.writeStream \
    .foreachBatch(idempotent_write) \
    .option("checkpointLocation", "/tmp/ckpt/orders") \
    .trigger(processingTime="10 seconds") \
    .start()

# Best practices for checkpoint setup

# 1. Use durable storage (S3, ADLS, GCS, HDFS) — NOT local disk
CHECKPOINT_DIR = "s3://my-bucket/spark-checkpoints/orders-pipeline"
# Or on Databricks DBFS:
# CHECKPOINT_DIR = "dbfs:/checkpoints/orders-pipeline"

# 2. Use fast storage for checkpoint (SSD-backed S3, not cold storage)

# 3. Never share checkpoint dirs between different queries

query = orders_df.writeStream \
    .format("delta") \
    .outputMode("append") \
    .option("checkpointLocation", CHECKPOINT_DIR) \
    .trigger(processingTime="30 seconds") \
    .start("s3://my-bucket/delta/orders")

# 4. To start FRESH (ignore previous checkpoint):
# Delete the checkpoint directory manually, then restart
# WARNING: this means reprocessing from the beginning!