Streaming: Kafka, Spark Streaming & Events

import json
from confluent_kafka import Producer

def delivery_callback(err, msg):
    """Callback triggered on broker ACK or failure."""
    if err is not None:
        print(f"Message delivery failed: {err}")
    else:
        print(f"Message delivered to {msg.topic()} [{msg.partition()}] at offset {msg.offset()}")

conf = {
    'bootstrap.servers': 'msk-cluster.aws.internal:9092',
    'client.id': 'market-data-producer',
    'acks': 'all',                          # Strong durability
    'enable.idempotence': True,             # Prevents duplicates on retry
    'compression.type': 'snappy',          # Low latency compression
    'linger.ms': 20,                        # Micro-batching for throughput
    'batch.num.messages': 5000,
}

producer = Producer(conf)

def publish_market_tick(ticker: str, price: float, volume: int):
    payload = {"ticker": ticker, "price": price, "volume": volume}
    # Key hashing on ticker ensures all ticks for 'AAPL' land on the exact same partition
    producer.produce(
        topic='market-ticks',
        key=ticker.encode('utf-8'),
        value=json.dumps(payload).encode('utf-8'),
        callback=delivery_callback
    )
    producer.poll(0)

# Flush remaining events on shutdown
producer.flush()

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StringType, DoubleType, LongType, TimestampType

spark = SparkSession.builder \
    .appName("KafkaMarketDataStreaming") \
    .config("spark.sql.streaming.stateStore.providerClass", 
            "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider") \
    .getOrCreate()

# Schema definition
tick_schema = StructType() \
    .add("ticker", StringType()) \
    .add("price", DoubleType()) \
    .add("volume", LongType()) \
    .add("trade_time", TimestampType())

# Read streaming feed from Kafka
kafka_stream = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "msk-cluster.aws.internal:9092") \
    .option("subscribe", "market-ticks") \
    .option("startingOffsets", "latest") \
    .load()

# Parse JSON payload
ticks_df = kafka_stream.select(
    F.from_json(F.col("value").cast("string"), tick_schema).alias("data")
).select("data.*")

# 5-minute tumbling candlestick aggregation with 10-minute watermark
ohlc_candles = ticks_df \
    .withWatermark("trade_time", "10 minutes") \
    .groupBy(
        F.window("trade_time", "5 minutes"),
        F.col("ticker")
    ) \
    .agg(
        F.first("price").alias("open"),
        F.max("price").alias("high"),
        F.min("price").alias("low"),
        F.last("price").alias("close"),
        F.sum("volume").alias("volume")
    )

# Write output to Delta Lake with checkpointing
query = ohlc_candles.writeStream \
    .format("delta") \
    .outputMode("append") \
    .option("checkpointLocation", "s3://lake-checkpoints/ohlc-candles/") \
    .start("s3://lake-silver/ohlc_candles_5min/")

query.awaitTermination()