Debugging Spark Streaming in a Notebook

Streams Never End

A streaming query doesn’t end. You call .start() and the query keeps running, waiting for new data forever. That’s the point in production, but in a notebook this means you either kill the cell by hand or you never really know if the stream finished processing the data you gave it.

I wanted to run a stream on a small slice of data, look at every micro-batch, and have the query stop on its own once there’s nothing left to process.

The Notebook Version

This is what I ended up with. It reads the change data feed of a Delta table, limits each micro-batch to a tiny size so I can see more than one batch, and polls the query status until the stream is idle.

import logging
import time

from pyspark.sql import DataFrame

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def batch_function(cdf_df: DataFrame, batch_id: int) -> None:
    logger.info("batch %s: %s rows", batch_id, cdf_df.count())


query = (
    spark.readStream
    .format("delta")
    .options(readChangeFeed="true", maxBytesPerTrigger="1K")
    .table("_temp_max_bytes_calc")
    .select("id", "year", "month")
    .drop_duplicates()
    .writeStream
    .foreachBatch(batch_function)
    .trigger(processingTime="5 seconds")
    .queryName("my query")
    .start()
)

time.sleep(5)
while query.isActive:
    stop_conditions = [
        not query.status["isDataAvailable"],
        not query.status["isTriggerActive"],
        query.status["message"] != "Initializing sources",
    ]

    if all(stop_conditions):
        query.stop()

    time.sleep(1)

query.awaitTermination(10)

A few of these lines need explaining.

• foreachBatch hands you each micro-batch as a regular DataFrame, along with its batch id. Inside the function you can do anything you’d do with a normal DataFrame, such as count it, display it, write it somewhere, etc. That makes foreachBatch the easiest place to see the stream’s work.
• maxBytesPerTrigger is a soft limit on how much data goes into one micro-batch. Setting the limit absurdly low (1K) forces the stream to split even a small table into several batches, so you can check that your logic also works across batches, since a single big batch can hide bugs. The source table needs change data feed enabled for readChangeFeed to work.
• query.status returns a dict with message, isDataAvailable and isTriggerActive. When there’s no data available, no trigger running, and the query is past its startup phase, the stream has caught up and is safe to stop. The "Initializing sources" check is there because right after start() the other two flags are also false, and without that check the loop would stop the query before Spark processed anything. The time.sleep(5) at the top is there for the same reason, and it’s doing more work than it looks. Before the sources even start initializing, the status message is "Initializing StreamExecution" with both flags false too, so without the sleep the loop could still stop the query too early.

The Simpler Way

The polling loop works, but Spark also has a method for exactly this, query.processAllAvailable(). Its docstring says it “blocks until all available data in the source has been processed and committed to the sink” and that it’s “intended for testing”. So the whole loop can be two lines.

query.processAllAvailable()
query.stop()

The catch from the docs is that if data keeps arriving, this can block forever, so use this on a fixed test table, not a live source.