Intermediate11 min read

Python and PySpark Coding Rounds

What actually gets asked in the coding round for a data engineer — file and JSON handling, idempotent transforms, and the PySpark DataFrame operations you should be able to write without thinking.

Data engineering coding rounds are usually not LeetCode-hard. They test whether you can handle messy data, write code that is safe to re-run, and use PySpark idiomatically instead of falling back to Python loops.

The three things they are really checking

Can you handle bad input? Missing keys, wrong types, nulls, duplicate records
Is your code idempotent? Running it twice must not double the data
Do you know the distributed way? Using DataFrame operations, not collect() then a Python loop

Python: parse and clean semi-structured records

A very common warm-up: *'here is a list of JSON events, some are malformed — clean them and aggregate.'* The graders look for defensive handling, not clever one-liners.

import json
from collections import defaultdict
from datetime import datetime


def parse_events(raw_lines):
    """Parse JSON lines, skipping bad records but counting them."""
    good, bad = [], 0
    for line in raw_lines:
        try:
            evt = json.loads(line)
        except json.JSONDecodeError:
            bad += 1
            continue

        # required fields — skip the record if any is missing
        site = evt.get("cell_site_id")
        ts = evt.get("event_ts")
        if site is None or ts is None:
            bad += 1
            continue

        try:
            evt["event_ts"] = datetime.fromisoformat(ts)
        except (TypeError, ValueError):
            bad += 1
            continue

        # volume may be missing -> treat as 0, not as a crash
        evt["volume_gb"] = float(evt.get("volume_gb") or 0)
        good.append(evt)

    return good, bad


def total_by_site(events):
    totals = defaultdict(float)
    for e in events:
        totals[e["cell_site_id"]] += e["volume_gb"]
    return dict(totals)

Tip — Always say what happens to bad records

Silently dropping rows is a red flag. Say out loud: 'In production I would route these to a quarantine table with the raw payload and the reason, and alert if the bad-record rate crosses a threshold.' That single sentence separates senior candidates from junior ones.

PySpark: the operations you must know cold

from pyspark.sql import SparkSession, Window
from pyspark.sql import functions as F

spark = SparkSession.builder.appName("kpi").getOrCreate()

df = spark.read.parquet("gs://bucket/raw/network_events/")

# --- filter + derive columns -------------------------------------------
df = (
    df.filter(F.col("event_date") == "2026-03-01")          # partition pruning
      .withColumn("volume_gb", F.col("volume_bytes") / 1e9)
      .withColumn("is_drop", F.col("call_status") == "DROPPED")
)

# --- aggregate ----------------------------------------------------------
kpi = (
    df.groupBy("cell_site_id", "event_date")
      .agg(
          F.sum("volume_gb").alias("total_gb"),
          F.count("*").alias("call_attempts"),
          F.sum(F.col("is_drop").cast("int")).alias("dropped_calls"),
          F.countDistinct("subscriber_id").alias("unique_subs"),
      )
      .withColumn(
          "drop_rate",
          F.round(F.col("dropped_calls") / F.col("call_attempts"), 4),
      )
)

# --- deduplicate: keep newest row per key -------------------------------
w = Window.partitionBy("subscriber_id").orderBy(F.col("updated_at").desc())
latest = (
    df.withColumn("rn", F.row_number().over(w))
      .filter(F.col("rn") == 1)
      .drop("rn")
)

# --- write idempotently: overwrite ONLY today's partition ---------------
(kpi.write
    .mode("overwrite")
    .option("partitionOverwriteMode", "dynamic")   # key line
    .partitionBy("event_date")
    .parquet("gs://bucket/core/network_kpi_daily/"))

Watch out — The idempotency line matters most

partitionOverwriteMode=dynamic replaces only the partitions present in the DataFrame, instead of wiping the whole table. Without it, a plain overwrite deletes every historical partition — a real production incident that interviewers love to probe.

Joins in PySpark

# Broadcast a SMALL dimension so no shuffle is needed on the big side
enriched = kpi.join(
    F.broadcast(dim_site),          # dim_site is a few thousand rows
    on="cell_site_id",
    how="left",
)

# Guard against fan-out: assert the dimension key is unique
assert dim_site.groupBy("cell_site_id").count().filter("count > 1").isEmpty()

Instead of...	Do this	Why
`df.collect()` then a Python loop	DataFrame transformations	`collect()` pulls all data to the driver and will OOM
A Python UDF	Built-in `F.*` functions, or `pandas_udf` if you must	Python UDFs serialise row-by-row and block Catalyst optimisation
`df.count()` inside a loop	Compute once, `cache()` if reused	Every action re-runs the whole lineage
`for` over partitions to write	`partitionBy(...)` on write	Spark parallelises the write for you

Practice questions

QWhat is the difference between a transformation and an action in Spark?show

Transformations (filter, select, join, groupBy) are lazy — they only build a plan. Actions (count, collect, show, write) trigger actual execution of that plan.

Laziness lets Catalyst optimise the whole chain — pushing filters down, pruning columns, combining steps
Practical consequence: nothing runs until an action, so a bug may only surface at write time
Calling multiple actions on the same DataFrame re-computes it unless you cache() or persist()

Likely follow-up: When is caching a bad idea?

QHow would you make a batch job safely re-runnable?show

Design every run so that repeating it for the same input produces the same result — no duplicates, no drift.

Overwrite by partition rather than append (partitionOverwriteMode=dynamic, or MERGE/DELETE+INSERT in BigQuery)
Derive the partition from the data, not from now() — so a rerun of a past date targets the right partition
Use deterministic keys — a surrogate key from a hash of the business key, not a random UUID
Make it stateless — the job should not depend on what the previous run left behind

QGiven a DataFrame of events, find each subscriber's first and last event in a single pass.show

Use a groupBy with min/max on the timestamp, or window functions if you also need the full row.

Simple: df.groupBy('subscriber_id').agg(F.min('event_ts'), F.max('event_ts'))
If you need whole rows: F.first(...)/F.last(...) with Window.partitionBy(...).orderBy(...) and an explicit frame
Mention that first/last without an orderBy are non-deterministic

QHow do you handle a schema change in an incoming file?show

Never let a source schema change silently break or silently pass. Detect it, then decide by change type.

Additive column — safe; land it in RAW, expose it downstream when needed
Removed or renamed column — fail loudly; downstream models depend on it
Type change — land as string in RAW, cast in STAGE so the raw record is never lost
Enforce with an explicit schema on read plus a schema test, not inferSchema

Likely follow-up: What is a data contract and how would you enforce one?

← PreviousSQL Interview Patterns You Must Know Next →Spark Architecture Explained Simply