How to use foreachBatch with a mock sink in PySpark
Demonstrates using Spark Structured Streaming's foreachBatch sink to capture and verify streaming batches by writing them into a custom mock sink object.
Requires third-party packages — install first
pip install pyspark
Python code
69 linesfrom pyspark.sql import SparkSession
from pyspark.sql.functions import col, lit
class MockSink:
def __init__(self):
self.batches = []
def write_batch(self, batch_df, batch_id):
# Collect batch data as list of dicts for verification
records = batch_df.collect()
self.batches.append({
"batch_id": batch_id,
"count": len(records),
"records": [row.asDict() for row in records]
})
def get_metrics(self):
return {
"total_batches": len(self.batches),
"total_records": sum(b["count"] for b in self.batches)
}
def main():
spark = SparkSession.builder.master("local[2]").appName("foreachBatchMockSink").getOrCreate()
spark.sparkContext.setLogLevel("ERROR")
# Create sample streaming DataFrame from static data (for demonstration)
sample_data = [
(1, "Alice", 25),
(2, "Bob", 30),
(3, "Charlie", 35),
(4, "Diana", 40)
]
source_df = spark.createDataFrame(sample_data, ["id", "name", "age"])
# Simulate streaming by using rate source with limited rows
streaming_df = spark.readStream.format("rate").option("rowsPerSecond", 2).load()
streaming_df = streaming_df.withColumn("id", col("value") + 1)
streaming_df = streaming_df.withColumn("name", lit("user"))
streaming_df = streaming_df.withColumn("age", col("value") * 10)
mock_sink = MockSink()
def write_to_mock(batch_df, batch_id):
mock_sink.write_batch(batch_df, batch_id)
query = streaming_df.writeStream \
.outputMode("append") \
.foreachBatch(write_to_mock) \
.trigger(processingTime="1 second") \
.start()
# Stop after 3 seconds to collect a few batches
import time
time.sleep(3)
query.stop()
metrics = mock_sink.get_metrics()
print(f"Total batches processed: {metrics['total_batches']}")
print(f"Total records processed: {metrics['total_records']}")
for batch in mock_sink.batches:
print(f" Batch {batch['batch_id']}: {batch['count']} records")
spark.stop()
if __name__ == "__main__":
main()Output
Total batches processed: 3
Total records processed: 6
Batch 0: 2 records
Batch 1: 2 records
Batch 2: 2 recordsHow it works
The foreachBatch sink lets you apply arbitrary logic to each micro-batch in a streaming query. Here the mock sink collects batch metadata and records as dicts for later verification. The rate source generates synthetic streaming data at 2 rows per second, and with a 1-second trigger, roughly 2 records arrive per batch. After 3 seconds, the query is stopped and metrics are printed to show how many batches and records were processed.
Common mistakes
- Forgetting to stop the query, leaving resources running indefinitely.
- Assuming batch processing time equals the trigger interval, while actual rows depend on the source's row rate.
- Calling `collect()` in a streaming write path can be memory-heavy for large batches; keep it only for small test data.
Variations
- Use `awaitTermination(timeout)` instead of `time.sleep` for cleaner shutdown handling.
- Replace the mock sink with a real data store like a Delta table or Kafka via `foreachBatch`.
Real-world use cases
- Validating streaming pipeline logic in unit tests by capturing and asserting on batch contents.
- Applying custom transformations or side effects to each micro-batch before writing to a target system.
- Implementing custom sinks for message queues or proprietary stores not natively supported by Spark.
Sponsored
More from Big data & Spark
Keep learning
Related tutorials and quizzes for this topic.