Two PySpark contracts Staff DEs learn on a bad Tuesday: foreachBatch restarts and applyInPandas skew

Your streaming MERGE looked fine in the notebook. Your applyInPandas transform looked elegant. Then a job restarted and silver row counts jumped, or a hot customer key OOMed a worker while the rest of the cluster sat idle.

Neither bug is exotic. Both are documented contracts that happy-path snippets skip.

This article covers two production surfaces Staff and Lead data engineers own on Databricks / PySpark:

1. Streaming upserts via foreachBatch: at-least-once writes, watermark dedupe state, Delta MERGE idempotency (logical, not txnAppId), empty batches, and why you must fully consume a batch DataFrame.
2. Pandas Function APIs: groupBy().applyInPandas() loads an entire group into memory and does not honor spark.sql.execution.arrow.maxRecordsPerBatch; mapInPandas / mapInArrow do.

I will mark FACT (from primary docs) vs interpretation (engineering judgment). No invented latency percentages. No "exactly-once foreachBatch" fairy tales.

Primary sources: Databricks foreachBatch, Delta MERGE / streaming upsert, watermarks, state store reader, idempotent Delta writes, pandas function APIs, Apache Arrow in PySpark.

---

Part 1. Design for the restart: foreachBatch, watermarks, and MERGE

Contract 1: foreachBatch is at-least-once

FACT (Databricks): foreachBatch() provides only at-least-once write guarantees. You can use the batchId to deduplicate and reason toward exactly-once yourself. You still own end-to-end semantics.

Interpretation: A green micro-batch in the UI is not a guarantee that a retry will be a no-op. Treat every handler as if Spark will call it again with the same batch_id.

APP_ID = "silver-orders-upsert-v1"  # change if you reset the checkpoint


def process_batch(batch_df, batch_id: int) -> None:
    # FACT: Delta sources can deliver empty micro-batches (e.g. OPTIMIZE with
    # nothing to process, or predicate pruning). Code must tolerate them.
    if batch_df.isEmpty():
        return

    valid = batch_df.filter("order_amount > 0 AND customer_id IS NOT NULL")
    invalid = batch_df.filter(
        "NOT (order_amount > 0 AND customer_id IS NOT NULL)"
    )

    # FACT: txnAppId + txnVersion on DataFrameWriter make appends idempotent
    # for retries of the same batch_id. Binding txnVersion to batch_id is the
    # documented pattern.
    (
        valid.write.format("delta")
        .mode("append")
        .option("txnAppId", APP_ID)
        .option("txnVersion", batch_id)
        .saveAsTable("main.silver.orders")
    )

    if not invalid.isEmpty():
        (
            invalid.write.format("delta")
            .mode("append")
            .option("txnAppId", APP_ID)
            .option("txnVersion", batch_id)
            .saveAsTable("main.silver.orders_dlq")
        )


(
    spark.readStream.table("main.bronze.raw_orders")
    .writeStream.foreachBatch(process_batch)
    .option("checkpointLocation", "/Volumes/main/ops/checkpoints/orders")
    .start()
)

FACT (Databricks Delta streaming): If you delete the streaming checkpoint and restart with a new checkpoint, you must provide a different txnAppId. New checkpoints start at batch ID 0; reusing the old app id can make Delta skip writes you intended to redo.

FACT (Databricks): Prefer letting the query fail and letting Lakeflow Jobs / Airflow retry the batch, rather than building silent local retry loops that can duplicate or drop data.

---

Contract 2: watermark dedupe has a state lease (and you must finish the batch)

FACT (Databricks watermarks): dropDuplicatesWithinWatermark always deduplicates events that arrive within the watermark delay. Events outside the delay might still be deduped, but that is not guaranteed. To guarantee dropping all duplicates, set the watermark delay greater than the maximum timestamp distance between duplicate events.

FACT (Databricks foreachBatch): When you use stateful operators such as dropDuplicatesWithinWatermark, each foreachBatch iteration must completely consume the batch DataFrame. If you only peek (show(2), take a few rows) and leave rows unread, the streaming query can fail on the next batch.

from pyspark.sql.functions import col

deduped = (
    spark.readStream.table("main.bronze.events")
    .withWatermark("event_time", "10 hours")  # > max duplicate lag (FACT guidance)
    .dropDuplicatesWithinWatermark(["event_id"])
)


def bad_handler(batch_df, batch_id: int) -> None:
    # Anti-pattern: partial consume after a stateful operator
    batch_df.show(2)
    # remaining rows never consumed -> next micro-batch can fail


def good_handler(batch_df, batch_id: int) -> None:
    if batch_df.isEmpty():
        return
    # Full action that consumes the plan (write, count, foreach, etc.)
    (
        batch_df.write.format("delta")
        .mode("append")
        .option("txnAppId", "events-dedupe-v1")
        .option("txnVersion", batch_id)
        .saveAsTable("main.silver.events")
    )


(
    deduped.writeStream.foreachBatch(good_handler)
    .option("checkpointLocation", "/Volumes/main/ops/checkpoints/events")
    .start()
)

Debugging tip (FACT): You can inspect Structured Streaming state with the statestore reader. Metadata can show an operator named dedupeWithinWatermark; query that operatorId when you need to see what the lease still holds. See Read Structured Streaming state information.

---

Contract 3: MERGE inside foreachBatch is a logical idempotency problem

FACT (Delta Lake): You must use foreachBatch for Delta Lake merge operations from Structured Streaming. Restarts can apply the same batch multiple times, so your merge statement inside foreachBatch must be idempotent.

FACT (Databricks / Delta writer options): txnAppId and txnVersion are DataFrameWriter options for idempotent writes (append / save). Do not assume they wrap DeltaTable.merge(...). Design MERGE so re-running the same micro-batch yields the same table state.

FACT (Delta): A merge can fail if multiple source rows match the same target row (ambiguous update). For CDC, preprocess to keep the latest change per key before merging.

FACT (Delta): MERGE often reads the source multiple times. That can inflate reported streaming input rates. Cache the micro-batch DataFrame before merge and unpersist after if this becomes a bottleneck.

from delta.tables import DeltaTable
from pyspark.sql import functions as F

target = DeltaTable.forName(spark, "main.silver.customers")


def upsert_cdc(batch_df, batch_id: int) -> None:
    if batch_df.isEmpty():
        return

    # Collapse to latest change per business key (FACT: avoid multi-match MERGE)
    latest = (
        batch_df.select(
            "customer_id",
            F.struct("event_time", "payload", "deleted").alias("chg"),
        )
        .groupBy("customer_id")
        .agg(F.max("chg").alias("chg"))
        .select("customer_id", "chg.*")
    )

    # FACT: MERGE may scan the source more than once -> cache the batch
    latest = latest.cache()
    try:
        (
            target.alias("t")
            .merge(latest.alias("s"), "t.customer_id = s.customer_id")
            .whenMatchedDelete(condition="s.deleted = true")
            .whenMatchedUpdate(
                set={"payload": "s.payload", "updated_at": "s.event_time"}
            )
            .whenNotMatchedInsert(
                condition="s.deleted = false",
                values={
                    "customer_id": "s.customer_id",
                    "payload": "s.payload",
                    "updated_at": "s.event_time",
                },
            )
            .execute()
        )
    finally:
        latest.unpersist()


(
    spark.readStream.table("main.bronze.customer_cdc")
    .writeStream.foreachBatch(upsert_cdc)
    .option("checkpointLocation", "/Volumes/main/ops/checkpoints/customer_cdc")
    .start()
)

Interpretation: "Idempotent MERGE" here means latest-wins on a stable key (and deletes keyed the same way), not "Spark will skip the merge via txn metadata." If the same CDC events replay, the table should converge to the same rows.

Managed alternative (context, not a product pitch): Databricks processing-guarantees docs note that managed Delta-to-Delta flows and AUTO CDC can give stronger defaults; a custom foreachBatch edge is still at-least-once unless you add explicit idempotent write or dedupe logic. Use AUTO CDC when it fits; when you hand-roll MERGE, you own the restart contract above.

---

Part 1 checklist (print this next to the PR)

Checkpoint	Why
Empty-batch safe	Delta sources can emit empty micro-batches
Full consume after stateful ops	Partial show/take can fail the next batch
Watermark > max duplicate lag	Out-of-watermark dedupe is not guaranteed
Append path uses txnAppId + txnVersion=batch_id	Writer-level idempotency on retry
New checkpoint => new txnAppId	Else batch 0 collisions skip writes
MERGE is latest-per-key / safe to re-run	txn options do not wrap DeltaTable.merge
Cache around MERGE when metrics / cost hurt	MERGE multi-pass over the micro-batch
Fail fast to the orchestrator	Local retry loops are how you invent exactly-once bugs

---

Part 2. Arrow is not a free pass: applyInPandas vs mapInPandas / mapInArrow

Python DE craft is not only streaming. The other Tuesday failure mode is "we used pandas because Arrow is fast" and then a skewed group killed the executor.

Three APIs, three memory contracts

FACT (Databricks pandas function APIs / Apache Spark):

API	What you write	Memory / batching contract
groupBy().applyInPandas(f, schema)	One pandas DataFrame per group	All rows for the group load into memory before f runs. spark.sql.execution.arrow.maxRecordsPerBatch is not applied. Skew => OOM risk.
mapInPandas(f, schema)	Iterator of pandas DataFrames (batches)	Batch size controlled by spark.sql.execution.arrow.maxRecordsPerBatch (default 10,000 rows in Spark docs). Output length can differ from input.
mapInArrow(f, schema)	Iterator of pyarrow.RecordBatch	Same Arrow batch sizing idea; stay in Arrow without pandas round-trip when you can.

Cogrouped applyInPandas has the same whole-cogroup-in-memory warning.

Skew turns "elegant" into OOM

# Looks clean. Under a hot customer_id it is a bomb.
def per_customer_features(pdf):
    # pdf = ALL rows for one customer_id
    pdf = pdf.sort_values("event_time")
    pdf["rolling_amt"] = pdf["amount"].rolling(50, min_periods=1).sum()
    return pdf


# FACT: entire group must fit in worker memory; maxRecordsPerBatch does NOT help
out = (
    events.groupBy("customer_id")
    .applyInPandas(
        per_customer_features,
        schema="customer_id string, event_time timestamp, amount double, rolling_amt double",
    )
)

Interpretation: Use applyInPandas when the group is a true unit of work and you can bound group size (filter, sample hot keys, pre-aggregate, salt, or split whales). Do not use it as a generic "run pandas on Spark" hammer.

Prefer batch maps when you do not need whole-group semantics

spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "5000")


def score_batches(iterator):
    # init once per Python worker task
    model = load_light_model()  # your code; keep it picklable / available on workers
    for pdf in iterator:
        pdf = pdf.copy()
        pdf["score"] = model.predict(pdf[["f1", "f2"]])
        yield pdf[["id", "score"]]


scored = events.mapInPandas(
    score_batches,
    schema="id long, score double",
)

Or stay in Arrow when pandas is unnecessary:

import pyarrow as pa
import pyarrow.compute as pc


def filter_arrow(iterator):
    for batch in iterator:
        mask = pc.greater(batch.column("amount"), 0)
        yield batch.filter(mask)


positive = events.mapInArrow(filter_arrow, schema=events.schema)

Decision table

Situation	Prefer	Watch out
Need whole-group window / asof / custom group ML	applyInPandas	Hot keys; no Arrow batch limit; cogroup same issue
Row/batch transform, arbitrary output length	mapInPandas	Tune maxRecordsPerBatch down if wide rows
CPU vector ops, avoid pandas copies	mapInArrow	Arrow type coverage / compute familiarity
Series-to-series column expression	classic pandas_udf	Output length must match input for many UDF shapes

FACT (Spark Arrow guide): Lower maxRecordsPerBatch when column count is large; Arrow batches can pressure JVM memory during conversion.

---

How the two themes connect

Both failures are the same Staff DE mistake: assuming the platform's happy path equals your failure domain.

• Streaming: you assumed merge + checkpoint = exactly-once. Docs say at-least-once + you make MERGE idempotent; writer txn options are a different tool.
• Pandas APIs: you assumed "Arrow optimized" = bounded memory. Docs say group maps ignore the batch limit.

Own the contract in the PR template, not in the postmortem.

---

Optional diagram (for a cover or inline figure)

Two panels:

1. Restart timeline: micro-batch N fails mid-MERGE -> orchestrator retries batch N -> arrows labeled "append: txn skip" vs "MERGE: must be latest-wins safe".
2. Skew funnel: many small customer_id groups sail through applyInPandas; one whale group expands to "all rows in RAM" with a red X; side path routes whales to salted / pre-aggregated / mapInPandas path.

---

References

1. Use foreachBatch to write to arbitrary data sinks (Databricks)
2. Delta Lake table deletes, updates, and merges (streaming foreachBatch section)
3. Apply watermarks / dropDuplicatesWithinWatermark (Databricks)
4. Read Structured Streaming state information
5. Delta Lake streaming reads and writes (idempotent foreachBatch writes)
6. pandas function APIs (Databricks)
7. Apache Arrow in PySpark
8. mapInArrow API

---

I'm Firat Celik, Lead Data Engineer. I write about production data platforms, PySpark/Delta, and the boring contracts that keep pipelines honest.