PySpark Interview Questions for Data Engineers — From Basics to Production

The Interviewer Says

We have two DataFrames, an orders DataFrame and a customers DataFrame. I want you to demonstrate different types of joins in PySpark, and then tell me how you would optimise a join when one of the DataFrames is very small compared to the other. Walk me through your approach.

Input Data

orders DataFrame

customers DataFrame

How to Handle This in the Interview

Before writing anything, say: Joins in PySpark work similarly to SQL joins but there are some important PySpark specific considerations around performance, especially around shuffle joins versus broadcast joins, and how to handle duplicate column names when both DataFrames share a join key with the same name.

Step 1: Know the join types available

• PySpark supports inner, left, right, full outer, left semi, left anti, and cross joins. Knowing when to use left semi and left anti is something that separates a good candidate from a great one.

Step 2: Handle the duplicate column name problem

• When both DataFrames have a column with the same name and you join using that column, PySpark keeps both copies in the result which creates ambiguous column references. Say: There are two ways to handle this, either drop the duplicate after joining or use the join condition as a string column name rather than a column expression, which automatically deduplicates the join key.

Step 3: Know when to use broadcast join

• Say: When one DataFrame is small enough to fit in memory on every executor, typically under a few hundred megabytes, a broadcast join is dramatically faster than a shuffle join because it eliminates the expensive shuffle operation entirely by sending the small DataFrame to every node.

Step 4: Mention the Catalyst optimiser

• Say: Spark's Catalyst optimiser will automatically apply a broadcast join if the smaller DataFrame is below the threshold set by spark.sql.autoBroadcastJoinThreshold, which defaults to 10 megabytes. For larger small tables I can force it manually using the broadcast hint.

Solution

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

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

# Create DataFrames
orders_data = [
    ("O01", "C01", "Laptop", 75000),
    ("O02", "C02", "Phone", 25000),
    ("O03", "C01", "Tablet", 30000),
    ("O04", "C03", "Monitor", 15000),
    ("O05", "C05", "Keyboard", 5000),
]
orders_df = spark.createDataFrame(
    orders_data,
    ["order_id", "customer_id", "product", "amount"]
)

customers_data = [
    ("C01", "Karan", "Mumbai"),
    ("C02", "Asha",  "Delhi"),
    ("C03", "Ravi",  "Pune"),
    ("C04", "Meena", "Chennai"),
]
customers_df = spark.createDataFrame(
    customers_data,
    ["customer_id", "name", "city"]
)

# INNER JOIN
# Returns only rows where customer_id exists in both DataFrames
inner_join = orders_df.join(customers_df, on="customer_id", how="inner")
print("INNER JOIN:")
inner_join.show()

# LEFT JOIN
# Returns all orders, NULL for customer details if no match
left_join = orders_df.join(customers_df, on="customer_id", how="left")
print("LEFT JOIN:")
left_join.show()

# RIGHT JOIN
# Returns all customers, NULL for order details if no match
right_join = orders_df.join(customers_df, on="customer_id", how="right")
print("RIGHT JOIN:")
right_join.show()

# FULL OUTER JOIN
# Returns everything from both sides
full_join = orders_df.join(customers_df, on="customer_id", how="full")
print("FULL OUTER JOIN:")
full_join.show()

# LEFT SEMI JOIN
# Returns only orders where a matching customer exists
# but does NOT include any columns from customers_df
semi_join = orders_df.join(customers_df, on="customer_id", how="left_semi")
print("LEFT SEMI JOIN:")
semi_join.show()

# LEFT ANTI JOIN
# Returns only orders where NO matching customer exists
# opposite of left semi join
anti_join = orders_df.join(customers_df, on="customer_id", how="left_anti")
print("LEFT ANTI JOIN:")
anti_join.show()

# BROADCAST JOIN
# Force broadcast when customers_df is small
# to avoid an expensive shuffle operation
broadcast_join = orders_df.join(
    F.broadcast(customers_df),
    on="customer_id",
    how="inner"
)
print("BROADCAST JOIN:")
broadcast_join.show()

# Handling duplicate column names when join key names differ
# or when you need both copies
orders_renamed = orders_df.withColumnRenamed("customer_id", "order_customer_id")
explicit_join = orders_renamed.join(
    customers_df,
    orders_renamed["order_customer_id"] == customers_df["customer_id"],
    how="inner"
)
explicit_join.show()

What to say: I pass the join key as a string when both DataFrames share the same column name, using on equals customer_id as a string, because PySpark then automatically keeps only one copy of the join key column in the result. If I passed it as a column expression like orders_df customer_id equals customers_df customer_id, PySpark would keep both copies and any subsequent reference to customer_id would be ambiguous and throw an AnalysisException. For the broadcast join I wrap the smaller DataFrame in F.broadcast to give Spark an explicit hint even if the DataFrame happens to be above the auto broadcast threshold.

---

Expected Outputs

INNER JOIN: 4 rows

| customer_id | order_id | product | amount | name | city |
|-------------|----------|----------|-------:|------|--------|
| C01 | O01 | Laptop | 75000 | Karan | Mumbai |
| C01 | O03 | Tablet | 30000 | Karan | Mumbai |
| C02 | O02 | Phone | 25000 | Asha | Delhi |
| C03 | O04 | Monitor | 15000 | Ravi | Pune |

LEFT JOIN: 5 rows (O05 kept with NULL customer details)

| customer_id | order_id | product | amount | name | city |
|-------------|----------|----------|-------:|------|--------|
| C01 | O01 | Laptop | 75000 | Karan | Mumbai |
| C01 | O03 | Tablet | 30000 | Karan | Mumbai |
| C02 | O02 | Phone | 25000 | Asha | Delhi |
| C03 | O04 | Monitor | 15000 | Ravi | Pune |
| C05 | O05 | Keyboard | 5000 | NULL | NULL |

LEFT SEMI JOIN: only order columns, matched rows only

| customer_id | order_id | product | amount |
|-------------|----------|----------|-------:|
| C01 | O01 | Laptop | 75000 |
| C01 | O03 | Tablet | 30000 |
| C02 | O02 | Phone | 25000 |
| C03 | O04 | Monitor | 15000 |

LEFT ANTI JOIN: only unmatched orders

| customer_id | order_id | product | amount |
|-------------|----------|----------|-------:|
| C05 | O05 | Keyboard | 5000 |

Follow-up Questions & Answers

Q1. What is the difference between left semi join and inner join? When would you use left semi join?

An inner join returns all columns from both DataFrames for matching rows, and if there are duplicate keys on the right side it multiplies the rows just like a SQL inner join. A left semi join returns only the columns from the left DataFrame for rows that have at least one match on the right, and it never duplicates rows regardless of how many matches exist on the right side. Left semi join is essentially a filter operation, it tells you which rows in the left DataFrame have a corresponding record on the right, without actually merging any columns from the right. Use left semi join when you want to filter one DataFrame based on the existence of a key in another, but do not need any data from the second DataFrame. A common real world example is filtering a large transactions table to only keep transactions whose customer exists in a valid customers table, without adding any customer columns to the result.

Q2. What is a shuffle join and why is it expensive?

A shuffle join, also called a sort merge join in Spark, is the default join strategy when both DataFrames are large. Spark needs to ensure all rows with the same join key end up on the same executor node before it can compare them. To achieve this it redistributes data across the network by hashing the join key and sending each row to the executor responsible for that hash bucket. This redistribution across the network is called a shuffle and it is one of the most expensive operations in Spark because it involves serialising data, writing it to disk, transferring it across the network between nodes, deserialising it on the receiving end, and then sorting it before the actual join can happen. On a cluster with hundreds of gigabytes of data, a shuffle can take minutes and is almost always the single biggest bottleneck in a Spark job.

Q3. How does a broadcast join avoid the shuffle and what are its limitations?

In a broadcast join, Spark takes the smaller DataFrame and sends a complete copy of it to every executor node in the cluster. Each executor then performs the join locally by looking up keys in its local copy of the small DataFrame, with no data movement between executors needed at all. This eliminates the shuffle entirely which is why it is so much faster. The limitation is size, the smaller DataFrame must be small enough to fit comfortably in the memory of every executor. If you broadcast a DataFrame that is too large, executors run out of memory and the job fails. The default auto broadcast threshold in Spark is 10 megabytes, controlled by spark.sql.autoBroadcastJoinThreshold. You can raise this if your executors have enough memory.

# Raise the auto broadcast threshold to 100MB
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", 100 * 1024 * 1024)

# Or force a broadcast hint manually regardless of size
result = large_df.join(F.broadcast(small_df), on="key", how="inner")

Q4. How would you handle the ambiguous column reference problem when both DataFrames have columns with the same name beyond just the join key?

When both DataFrames have multiple columns with the same name, joining them creates ambiguity for every shared column not just the key. The cleanest solutions are to rename columns before joining or to select and alias after joining.

# Option 1: Rename conflicting columns before joining
orders_df2 = orders_df.withColumnRenamed("amount", "order_amount")
result = orders_df2.join(customers_df, on="customer_id", how="inner")

# Option 2: Use DataFrame specific column references after joining
# when both have a column called 'amount'
result = orders_df.join(
    customers_df,
    on=orders_df["customer_id"] == customers_df["customer_id"],
    how="inner"
).select(
    orders_df["customer_id"],
    orders_df["order_id"],
    orders_df["amount"].alias("order_amount"),
    customers_df["name"],
    customers_df["city"]
)

# Option 3: Add a prefix to all columns of one DataFrame before joining
from functools import reduce
prefixed_customers = customers_df.select(
    [F.col(c).alias("cust_" + c) for c in customers_df.columns]
)
result = orders_df.join(
    prefixed_customers,
    orders_df["customer_id"] == prefixed_customers["cust_customer_id"],
    how="inner"
)

What to say: Option 1 is the cleanest for known column conflicts. Option 2 is useful when you need precise control over exactly which columns appear in the result. Option 3, prefixing all columns from one side, is a useful defensive pattern in generic pipeline code where you cannot predict which column names might clash, since it guarantees no ambiguity regardless of the source schema.

Q5. What would you check in the Spark UI if a join is running much slower than expected?

What to say: The Spark UI is the first place I go when a job is slower than expected. For a slow join I look at three specific things.

First I check the DAG in the SQL tab to confirm which join strategy Spark actually chose. If I expected a broadcast join but the plan shows a SortMergeJoin, it means the DataFrame was above the broadcast threshold and I need to either raise the threshold or explicitly add the broadcast hint.

Second I check the Stages tab for any stage with an unusually high amount of data read or written. A shuffle join stage that reads and writes hundreds of gigabytes signals that the shuffle itself is the bottleneck, and I would consider repartitioning the DataFrames on the join key before the join to reduce the shuffle cost.

Third I check for data skew in the Tasks tab within the shuffle stage. If most tasks finish in seconds but a handful take minutes, that means one or a few join key values have a disproportionately large number of rows, causing certain executors to receive far more data than others. This is called data skew and it requires a different fix, such as salting the join key to spread the skewed data across multiple partitions.

# Salting technique for skewed joins
import random

# Add a random salt to the skewed key in the large DataFrame
orders_salted = orders_df.withColumn(
    "salted_key",
    F.concat(F.col("customer_id"), F.lit("_"), (F.rand() * 10).cast("int"))
)

# Explode the small DataFrame to match all possible salt values
from pyspark.sql.functions import array, explode, lit
customers_salted = customers_df.withColumn(
    "salt",
    explode(array([lit(i) for i in range(10)]))
).withColumn(
    "salted_key",
    F.concat(F.col("customer_id"), F.lit("_"), F.col("salt"))
)

# Join on salted key instead of raw customer_id
result = orders_salted.join(
    customers_salted,
    on="salted_key",
    how="inner"
).drop("salted_key", "salt")

What to say: Salting works by artificially distributing the rows of a skewed key across multiple buckets by appending a random number to the key. The small DataFrame is expanded to have matching rows for every possible salt value. After joining on the salted key the salt columns are dropped. This spreads what was one overloaded task into ten smaller tasks that run in parallel, eliminating the straggler executor problem that skew causes.

The Interviewer Says

We have a sales DataFrame with transactions across multiple products and regions. I want you to first perform some GroupBy aggregations to get summary statistics, and then use window functions to add running totals and rankings within groups, all without losing the row level detail. Walk me through how you think about when to use GroupBy versus window functions.

Input Data

How to Handle This in the Interview

Before writing anything, say: GroupBy collapses multiple rows into one summary row per group, which means you lose the row level detail. Window functions add computed values alongside each row without collapsing anything. The choice between them depends entirely on whether the business needs a summary or a row level enrichment.

Step 1: GroupBy for summary statistics

• GroupBy with agg is the right choice when the output should be one row per group, for example total sales per product or average amount per region. The original rows are gone from the result.

Step 2: Window functions for row level enrichment

• When the business wants to see the running total, rank, or moving average alongside each individual transaction, use window functions. The original rows are preserved and the computed value is added as a new column.

Step 3: Know the Window specification components

• A Window spec has three optional parts, partitionBy which defines the group boundary, orderBy which defines the sort order within the group, and rowsBetween or rangeBetween which defines the frame. Knowing when each part is needed and what happens when one is omitted is important.

Step 4: Know which functions need an orderBy

• Ranking functions like rank, dense_rank, row_number, and lag always need an orderBy. Aggregation window functions like sum and avg work with or without orderBy but the frame changes depending on whether orderBy is present.

Solution

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

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

data = [
    ("T01", "P01", "North", "2024-01-05", 5000),
    ("T02", "P01", "South", "2024-01-06", 7000),
    ("T03", "P02", "North", "2024-01-07", 3000),
    ("T04", "P01", "North", "2024-01-08", 6000),
    ("T05", "P02", "South", "2024-01-09", 4000),
    ("T06", "P02", "North", "2024-01-10", 8000),
    ("T07", "P01", "South", "2024-01-11", 9000),
    ("T08", "P02", "South", "2024-01-12", 2000),
]

df = spark.createDataFrame(
    data,
    ["txn_id", "product_id", "region", "sales_date", "amount"]
)

# -------------------------------------------------------
# PART 1: GroupBy Aggregations
# -------------------------------------------------------

# Total sales, average amount, max amount, and count per product
product_summary = df.groupBy("product_id").agg(
    F.sum("amount").alias("total_sales"),
    F.avg("amount").alias("avg_sales"),
    F.max("amount").alias("max_sales"),
    F.count("*").alias("txn_count")
)
print("Product Summary:")
product_summary.show()

# Total sales per product per region
product_region_summary = df.groupBy("product_id", "region").agg(
    F.sum("amount").alias("total_sales")
).orderBy("product_id", "region")
print("Product Region Summary:")
product_region_summary.show()

# Collect all transaction amounts per product as a list
product_txn_list = df.groupBy("product_id").agg(
    F.collect_list("amount").alias("all_amounts"),
    F.collect_set("region").alias("unique_regions")
)
print("Product Transaction List:")
product_txn_list.show(truncate=False)

# -------------------------------------------------------
# PART 2: Window Functions
# -------------------------------------------------------

# Window spec: partition by product, order by sales_date
w_product = Window.partitionBy("product_id").orderBy("sales_date")

# Window spec: partition by product, full partition for aggregations
w_product_full = Window.partitionBy("product_id")

# Running total of amount per product ordered by date
df_with_running_total = df.withColumn(
    "running_total",
    F.sum("amount").over(
        Window.partitionBy("product_id")
              .orderBy("sales_date")
              .rowsBetween(Window.unboundedPreceding, Window.currentRow)
    )
)
print("Running Total per Product:")
df_with_running_total.show()

# Rank each transaction within its product by amount descending
df_with_rank = df.withColumn(
    "rank",
    F.rank().over(
        Window.partitionBy("product_id")
              .orderBy(F.col("amount").desc())
    )
).withColumn(
    "dense_rank",
    F.dense_rank().over(
        Window.partitionBy("product_id")
              .orderBy(F.col("amount").desc())
    )
).withColumn(
    "row_number",
    F.row_number().over(
        Window.partitionBy("product_id")
              .orderBy(F.col("amount").desc())
    )
)
print("Ranks per Product:")
df_with_rank.show()

# LAG and LEAD: previous and next transaction amount per product
df_with_lag_lead = df.withColumn(
    "prev_amount",
    F.lag("amount", 1).over(w_product)
).withColumn(
    "next_amount",
    F.lead("amount", 1).over(w_product)
).withColumn(
    "amount_diff_from_prev",
    F.col("amount") - F.lag("amount", 1).over(w_product)
)
print("LAG and LEAD:")
df_with_lag_lead.show()

# Total sales of the entire product group attached to each row
df_with_group_total = df.withColumn(
    "product_total",
    F.sum("amount").over(w_product_full)
).withColumn(
    "pct_of_product_total",
    F.round(F.col("amount") * 100 / F.sum("amount").over(w_product_full), 2)
)
print("Each transaction as percentage of product total:")
df_with_group_total.show()

What to say: I define the Window specification once and reuse it across multiple withColumn calls rather than repeating the full window definition each time, which keeps the code readable and avoids any chance of inconsistency between window specs that should be identical. The key distinction I always make clear is that groupBy produces one row per group while window functions keep every original row, and that difference drives which one I reach for based on what the output needs to look like.

Expected Output: Running Total

Follow-up Questions & Answers

Q1. What is the difference between groupBy with agg and a window function using sum over, when should you use each?

groupBy with agg collapses the DataFrame to one row per group and computes the aggregate across all rows in that group. The individual transaction rows are gone from the result, replaced by a single summary row. A window function using sum over keeps every row and attaches the aggregate as a new column on each row. So if you have ten transactions for product P01 and you want one total, use groupBy. If you want each of those ten transactions to show both its own amount and the running total so far, use a window function. A practical way to remember this is to ask whether the business question is about the group as a whole, use groupBy, or about each individual row in the context of its group, use window function.

Q2. What is the difference between rowsBetween and rangeBetween in a window frame?

Both define which rows are included in the window calculation relative to the current row, but they differ in how they interpret the boundaries.

# rowsBetween uses physical row offsets
# exactly 3 rows before and including the current row
w_rows = Window.partitionBy("product_id") \
               .orderBy("sales_date") \
               .rowsBetween(-3, Window.currentRow)

# rangeBetween uses value based offsets on the ORDER BY column
# all rows whose sales_date is within 3 days before the current row
w_range = Window.partitionBy("product_id") \
                .orderBy("sales_date") \
                .rangeBetween(-3, Window.currentRow)

# Common constants
# Window.unboundedPreceding = from the very first row of the partition
# Window.currentRow         = up to and including the current row
# Window.unboundedFollowing = up to and including the very last row

What to say: rowsBetween counts actual physical rows regardless of the values in the ORDER BY column, so rowsBetween minus 2 and currentRow always includes exactly three rows. rangeBetween interprets the offset as a value range on the ORDER BY column, so for a date column rangeBetween minus 3 and currentRow includes all rows whose date is within three days before the current row's date, which could be zero rows or ten rows depending on the data. For running totals and rankings I always use rowsBetween since it gives predictable, deterministic behaviour. rangeBetween is useful when the frame should be defined by a time or value interval rather than a fixed row count.

Q3. What happens if you define a Window with partitionBy but no orderBy, and then use rank on it?

This throws an AnalysisException at runtime because ranking functions require an ordering to determine positions. There is no meaningful concept of which row is first or second without knowing the sort order, so Spark refuses to run it.

# This throws AnalysisException
w_no_order = Window.partitionBy("product_id")
df.withColumn("rank", F.rank().over(w_no_order))  # ERROR

# This works because orderBy is specified
w_with_order = Window.partitionBy("product_id").orderBy("amount")
df.withColumn("rank", F.rank().over(w_with_order))  # OK

What to say: Ranking functions like rank, dense_rank, row_number, lag, and lead always need an orderBy inside the window spec. Aggregation functions like sum, avg, min, and max can work with or without orderBy, but the behaviour changes significantly. Without orderBy, sum over the full partition gives the grand total of the group for every row. With orderBy, sum defaults to a cumulative running total up to the current row. This difference catches many people off guard and is a common source of silent incorrect results.

Q4. How would you find the top one transaction per product by amount using window functions without using groupBy?

from pyspark.sql.window import Window
from pyspark.sql import functions as F

w_rank = Window.partitionBy("product_id").orderBy(F.col("amount").desc())

top_per_product = df.withColumn(
    "row_num",
    F.row_number().over(w_rank)
).filter(F.col("row_num") == 1) \
 .drop("row_num")

top_per_product.show()

What to say: This is a very common interview pattern. I use row_number rather than rank here specifically because row_number assigns a unique number to every row even when amounts are tied, guaranteeing exactly one row per product in the output. If I used rank and two transactions in P01 had the same highest amount, both would get rank one and I would get two rows for P01. Whether tied rows should both appear or only one should be picked is exactly the question I would ask the interviewer before choosing between row_number and rank.

Q5. What is the performance implication of having too many window functions with different partition and order specifications in the same query?

Each distinct window specification triggers a separate shuffle in Spark since data needs to be repartitioned and sorted differently for each unique combination of partitionBy and orderBy columns. If you have five window functions but three of them share the same window spec and two use a different one, Spark performs two separate shuffles.

# Less efficient: three different window specs trigger three shuffles
df.withColumn("col1", F.sum("amount").over(Window.partitionBy("product_id").orderBy("sales_date"))) \
  .withColumn("col2", F.rank().over(Window.partitionBy("product_id").orderBy("amount"))) \
  .withColumn("col3", F.sum("amount").over(Window.partitionBy("region").orderBy("sales_date")))

# More efficient: reuse the same window spec object wherever possible
w1 = Window.partitionBy("product_id").orderBy("sales_date")
w2 = Window.partitionBy("product_id").orderBy("amount")
w3 = Window.partitionBy("region").orderBy("sales_date")

df.withColumn("running_total", F.sum("amount").over(w1)) \
  .withColumn("amount_rank",   F.rank().over(w2)) \
  .withColumn("region_total",  F.sum("amount").over(w3))

What to say: Defining window specs as named variables and reusing them is both more readable and more efficient. Spark recognises when multiple withColumn calls use the same window spec and combines them into a single shuffle and sort rather than repeating the expensive redistribution multiple times. I always extract window specs into named variables at the top of a transformation block and reuse them consistently, which is a small habit that makes a meaningful difference in job performance when the DataFrame is large.

The Interviewer Says

In real world data pipelines, null values are one of the most common and dangerous problems you will face. A null in the wrong place can silently corrupt aggregations, break joins, or cause entire pipeline stages to fail. I have a customer transactions DataFrame with several null values scattered across different columns. Walk me through how you would detect, analyse, and handle nulls in PySpark, and tell me the different strategies you would consider for each type of column.

Input Data

How to Handle This in the Interview

Before writing anything, say: Null handling is not a one size fits all operation. Before deciding what to do with nulls I always want to first understand how many there are, which columns have them, and what the business meaning of a null is in each column. A null customer_id probably means a data quality failure and that row should be dropped. A null amount might mean the transaction is pending and should be filled with zero or flagged separately. A null region might mean the customer is international and can be filled with a default value.

Step 1: Quantify nulls before doing anything

• Say: The first thing I always do is write a null count check across every column so I know the scale of the problem before choosing a strategy. This is a diagnostic step, not a transformation step.

Step 2: Understand the three strategies

• Drop rows with nulls using dropna when a null makes the entire row meaningless. Fill nulls with a default value using fillna when a sensible default exists. Replace nulls with a derived value like a column mean or median when dropping or using a fixed default would distort the analysis.

Step 3: Apply different strategies per column

• Never apply the same null handling strategy to all columns blindly. Different columns have different business meanings and different appropriate responses to a null value.

Step 4: Null safe comparisons

• Mention that normal equality comparisons with null in PySpark behave the same as SQL, null equals anything is always false. Show the null safe equality operator or IS NULL filter for detecting nulls.

Solution

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import (
    StructType, StructField,
    StringType, IntegerType, DateType
)

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

data = [
    ("T01", "C01", "Laptop", 75000, "2024-01-05", "North"),
    ("T02", "C02", None, 25000, "2024-01-06", "South"),
    ("T03", None, "Tablet", None, "2024-01-07", "North"),
    ("T04", "C03", "Monitor", 15000, None, None),
    ("T05", "C04", "Keyboard", 5000, "2024-01-09", "South"),
    ("T06", "C05", "Phone", None, "2024-01-10", None),
]

df = spark.createDataFrame(
    data,
    ["txn_id", "customer_id", "product", "amount", "txn_date", "region"]
)

print("Raw DataFrame:")
df.show()

# -------------------------------------------------------
# STEP 1: Quantify nulls across all columns
# -------------------------------------------------------

print("Null counts per column:")
null_counts = df.select([
    F.count(
        F.when(F.col(c).isNull(), c)
    ).alias(c)
    for c in df.columns
])
null_counts.show()

# Null count as percentage of total rows
total_rows = df.count()
print(f"Total rows: {total_rows}")

null_pct = df.select([
    F.round(
        F.count(F.when(F.col(c).isNull(), c)) * 100.0 / total_rows,
        2
    ).alias(c + "_null_pct")
    for c in df.columns
])
null_pct.show()

# -------------------------------------------------------
# STEP 2: dropna strategies
# -------------------------------------------------------

# Drop rows where ANY column is null
df_drop_any = df.dropna(how="any")
print("After dropping rows with ANY null:")
df_drop_any.show()

# Drop rows where ALL columns are null
df_drop_all = df.dropna(how="all")
print("After dropping rows where ALL columns are null:")
df_drop_all.show()

# Drop rows where specific critical columns are null
# customer_id is critical, a transaction with no customer is meaningless
df_drop_critical = df.dropna(subset=["customer_id"])
print("After dropping rows where customer_id is null:")
df_drop_critical.show()

# -------------------------------------------------------
# STEP 3: fillna strategies
# -------------------------------------------------------

# Fill nulls with fixed values per column
df_filled = df.fillna({
    "product": "Unknown",
    "region": "International",
    "amount": 0
})
print("After fillna with fixed values:")
df_filled.show()

# Fill null amount with the mean amount
mean_amount = df.select(F.mean("amount")).collect()[0][0]
print(f"Mean amount: {mean_amount}")

df_filled_mean = df.withColumn(
    "amount",
    F.when(
        F.col("amount").isNull(),
        F.lit(round(mean_amount, 2))
    ).otherwise(F.col("amount"))
)
print("After filling null amount with mean:")
df_filled_mean.show()

# -------------------------------------------------------
# STEP 4: Null safe operations
# -------------------------------------------------------

# Normal equality fails with nulls
# col("region") == "North" returns false for null rows, not true or false
df_north = df.filter(F.col("region") == "North")
print("Filter region == North (null rows excluded silently):")
df_north.show()

# Explicit null check
df_null_region = df.filter(F.col("region").isNull())
print("Rows where region is null:")
df_null_region.show()

# Null safe equality using eqNullSafe
df_null_safe = df.filter(F.col("region").eqNullSafe("North"))
print("Null safe equality for North:")
df_null_safe.show()

# -------------------------------------------------------
# STEP 5: Replace nulls with a window derived value
# -------------------------------------------------------
from pyspark.sql.window import Window

# Add region to the data to make this meaningful
# Fill null amount with the average amount of the same region
w_region = Window.partitionBy("region")

df_with_region_avg = df.withColumn(
    "region_avg_amount",
    F.avg("amount").over(w_region)
).withColumn(
    "amount_cleaned",
    F.when(
        F.col("amount").isNull(),
        F.round(F.col("region_avg_amount"), 2)
    ).otherwise(F.col("amount"))
).drop("region_avg_amount")

print("Amount filled with region average:")
df_with_region_avg.show()

What to say: I always start with the null count diagnostic before touching any data because the right strategy depends entirely on how many nulls there are and in which columns. Dropping rows is destructive and should only be done for columns where a null makes the entire row invalid, like a transaction with no customer_id. Filling with a fixed value like zero or Unknown is appropriate for optional fields where a default is meaningful. Filling with a computed value like the column mean or a group average is appropriate for numerical fields where you want to preserve the statistical distribution of the data rather than distorting it with zeros.

Expected Output — Null Counts

Follow-up Questions & Answers

Q1. What is the difference between dropna and filter with isNull? When would you use each?

dropna is a high level convenience method specifically designed for dropping rows based on null conditions across one or more columns. filter with isNull is a general purpose row filter that gives more precise control.

# dropna: drop rows where amount OR region is null
df.dropna(subset=["amount", "region"])

# Equivalent using filter but more explicit
df.filter(F.col("amount").isNotNull() & F.col("region").isNotNull())

# dropna cannot express this: keep rows where amount is null
# but only if region is also null
# for this you need filter
df.filter(
    F.col("amount").isNotNull() |
    (F.col("amount").isNull() & F.col("region").isNotNull())
)

What to say: dropna is faster to write for common patterns like drop any row missing a critical column. filter with isNull and isNotNull gives you full boolean logic for complex multi column null conditions that dropna cannot express. In production pipelines I tend to use filter for null handling because the explicit boolean logic makes the intent obvious to anyone reading the code later, whereas dropna hides the column list inside a method argument that is easy to miss during a code review.

Q2. What happens to null values in a groupBy aggregation like sum or count?

This is one of the most important null behaviours to understand in Spark because it matches SQL behaviour and can silently produce wrong results if you are not aware of it.

# SUM ignores nulls, so null amounts are simply skipped
df.groupBy("region").agg(F.sum("amount").alias("total"))
# A region with 3 transactions where 2 have null amounts
# will show only the sum of the 1 non-null amount

# COUNT star counts all rows including those with nulls
df.groupBy("region").agg(F.count("*").alias("row_count"))

# COUNT on a specific column counts only non-null values
df.groupBy("region").agg(F.count("amount").alias("non_null_count"))

# To explicitly include nulls in count
df.groupBy("region").agg(
    F.count("*").alias("total_rows"),
    F.count("amount").alias("non_null_amounts"),
    (F.count("*") - F.count("amount")).alias("null_amounts")
)

What to say: SUM, AVG, MIN, and MAX all silently skip null values, which means the result is computed only over the non null rows. This is usually the correct behaviour but it can be misleading if a large fraction of values are null, since the average of ten values where three are null is computed over only seven rows, not ten, and the denominator is silently different from what you might expect. I always check null counts before running aggregations on numerical columns to understand how representative the result actually is.

Q3. How does coalesce work in PySpark and how is it different from fillna?

Both replace nulls but they operate differently and at different levels of flexibility.

# coalesce takes multiple column expressions and returns the first non-null one
# useful when you want to fall back through a priority list of columns
df.withColumn(
    "effective_region",
    F.coalesce(
        F.col("region"),
        F.col("customer_region"),  # fall back to another column
        F.lit("Unknown")           # final default if all are null
    )
)

# fillna replaces nulls with a fixed scalar value per column
# cannot reference other columns as the replacement value
df.fillna({"region": "Unknown", "amount": 0})

# coalesce is more powerful when the replacement value is dynamic
# for example fill null amount with the previous row's amount
from pyspark.sql.window import Window
w = Window.partitionBy("customer_id").orderBy("txn_date")
df.withColumn(
    "amount_filled",
    F.coalesce(
        F.col("amount"),
        F.lag("amount", 1).over(w)  # use previous row's amount if current is null
    )
)

What to say: fillna is convenient for filling nulls with simple fixed scalar values across one or more columns in one call. coalesce is more powerful because the replacement value can itself be a column expression, another column's value, a window function result, or any computed expression. I reach for fillna for quick fixed value fills and coalesce whenever the replacement logic is conditional or references another column.

Q4. How would you write a reusable null audit function that works on any DataFrame regardless of its schema?

This is a very practical production skill since null auditing should run on every DataFrame entering a pipeline.

def null_audit(df, spark):
    total = df.count()
    audit_rows = []

    for col_name in df.columns:
        null_count = df.filter(F.col(col_name).isNull()).count()
        null_pct = round(null_count * 100.0 / total, 2) if total > 0 else 0.0
        dtype = dict(df.dtypes)[col_name]
        audit_rows.append((col_name, dtype, null_count, null_pct))

    audit_df = spark.createDataFrame(
        audit_rows,
        ["column_name", "data_type", "null_count", "null_pct"]
    )
    return audit_df.orderBy(F.col("null_count").desc())

audit_result = null_audit(df, spark)
audit_result.show()

What to say: I use df.dtypes which returns a list of column name and type tuples to make the function fully schema agnostic. It works on any DataFrame without needing to know the column names in advance. In a real pipeline I would call this function at the entry point of every transformation stage and write the results to a data quality log table with a timestamp and the source table name, so there is a historical record of null rates over time that can be tracked and alerted on.

Q5. What is the difference between null and NaN in PySpark, and why does it matter?

This is a subtle but important distinction that catches many people off guard.

import math
from pyspark.sql.types import DoubleType

# NaN (Not a Number) is a special floating point value
# it is NOT the same as null
# NaN appears when you do things like 0.0 / 0.0 in floating point arithmetic

data_with_nan = [
    ("T01", 5000.0),
    ("T02", float("nan")),  # NaN
    ("T03", None),           # null
]
df_nan = spark.createDataFrame(data_with_nan, ["txn_id", "amount"])

# isNull does NOT catch NaN
df_nan.filter(F.col("amount").isNull()).show()
# Returns only T03

# isnan catches NaN but NOT null
df_nan.filter(F.isnan("amount")).show()
# Returns only T02

# To catch both null and NaN together
df_nan.filter(
    F.col("amount").isNull() | F.isnan("amount")
).show()
# Returns both T02 and T03

# Replace both null and NaN with zero
df_cleaned = df_nan.withColumn(
    "amount",
    F.when(
        F.col("amount").isNull() | F.isnan("amount"),
        F.lit(0.0)
    ).otherwise(F.col("amount"))
)
df_cleaned.show()

What to say: NaN is a IEEE 754 floating point concept meaning the result of an undefined arithmetic operation, like dividing zero by zero. Null means the value is absent or unknown. In PySpark these are two completely separate things that require different detection methods. isNull only catches null and isnan only catches NaN. If you fill nulls with fillna but have NaN values in the same column, the NaN values pass through untouched and can silently corrupt downstream aggregations since most Spark aggregation functions propagate NaN, meaning a single NaN in a sum produces NaN as the total. I always check for both null and NaN in numerical columns when doing data quality validation.

The Interviewer Says

In a real data engineering pipeline you will rarely work with just CSV files. You need to know how to read and write different file formats efficiently, understand the trade-offs between them, and know when to use which format. Walk me through how you would read and write CSV, JSON, Parquet, and Delta formats in PySpark, and explain when you would choose each one.

How to Handle This in the Interview

Before writing anything, say: File format choice is one of the most impactful decisions in a data pipeline because it directly affects read performance, storage cost, schema evolution capability, and whether downstream consumers can query the data efficiently. Let me walk through each format and its real world use case.

Step 1: Know the core trade-offs

• CSV is human readable and universally supported but has no schema, no compression by default, and no support for nested data. JSON supports nested structures but is verbose and slow to parse at scale. Parquet is columnar, compressed, schema aware, and ideal for analytical workloads. Delta adds ACID transactions, time travel, and schema enforcement on top of Parquet.

Step 2: Know which format to recommend in an interview

• Say: For any serious data engineering workload I would always recommend Parquet as the minimum and Delta Lake if the platform supports it. CSV and JSON are acceptable only at the ingestion boundary where you have no control over what the upstream system sends.

Step 3: Understand partition pruning and predicate pushdown

• Writing data partitioned by a column like date or region allows Spark to skip entire folders when reading with a filter on that column, which dramatically reduces the amount of data scanned. This is called partition pruning and it only works with columnar formats like Parquet and Delta.

Step 4: Mention schema evolution

• Ask: Does the schema of this data change over time? If yes, I would lean toward Delta which handles schema evolution cleanly with mergeSchema and enforceSchema options.

Solution

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import (
    StructType, StructField,
    StringType, IntegerType, DateType
)

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

data = [
    ("T01", "C01", "Laptop", 75000, "2024-01-05", "North"),
    ("T02", "C02", "Phone", 25000, "2024-01-06", "South"),
    ("T03", "C03", "Tablet", 30000, "2024-01-07", "North"),
    ("T04", "C04", "Monitor", 15000, "2024-01-08", "South"),
    ("T05", "C05", "Keyboard", 5000, "2024-01-09", "North"),
]

df = spark.createDataFrame(
    data,
    ["txn_id", "customer_id", "product", "amount", "txn_date", "region"]
)

# -------------------------------------------------------
# CSV: Read and Write
# -------------------------------------------------------

# Write CSV
df.write \
    .option("header", "true") \
    .mode("overwrite") \
    .csv("/output/transactions_csv")

# Read CSV with explicit schema
schema = StructType([
    StructField("txn_id", StringType(), True),
    StructField("customer_id", StringType(), True),
    StructField("product", StringType(), True),
    StructField("amount", IntegerType(), True),
    StructField("txn_date", StringType(), True),
    StructField("region", StringType(), True)
])

df_csv = spark.read \
    .option("header", "true") \
    .option("mode", "DROPMALFORMED") \
    .schema(schema) \
    .csv("/output/transactions_csv")

print("CSV Read:")
df_csv.show()

# -------------------------------------------------------
# JSON: Read and Write
# -------------------------------------------------------

# Write JSON
df.write \
    .mode("overwrite") \
    .json("/output/transactions_json")

# Read JSON with schema inference is acceptable since JSON carries type hints
df_json = spark.read \
    .option("mode", "PERMISSIVE") \
    .schema(schema) \
    .json("/output/transactions_json")

print("JSON Read:")
df_json.show()

# -------------------------------------------------------
# Parquet: Read and Write
# -------------------------------------------------------

# Write Parquet, partitioned by region
df.write \
    .mode("overwrite") \
    .partitionBy("region") \
    .parquet("/output/transactions_parquet")

# Read all Parquet data
df_parquet = spark.read.parquet("/output/transactions_parquet")
print("Parquet Read (all):")
df_parquet.show()

# Read with partition pruning, only North region is scanned
df_parquet_north = spark.read \
    .parquet("/output/transactions_parquet") \
    .filter(F.col("region") == "North")
print("Parquet Read (North only, partition pruning):")
df_parquet_north.show()

# -------------------------------------------------------
# Delta Lake: Read and Write
# -------------------------------------------------------

# Write Delta
df.write \
    .format("delta") \
    .mode("overwrite") \
    .partitionBy("region") \
    .save("/output/transactions_delta")

# Read Delta
df_delta = spark.read \
    .format("delta") \
    .load("/output/transactions_delta")
print("Delta Read:")
df_delta.show()

# Delta Time Travel: read a previous version
df_delta_v0 = spark.read \
    .format("delta") \
    .option("versionAsOf", 0) \
    .load("/output/transactions_delta")
print("Delta Time Travel (version 0):")
df_delta_v0.show()

# Delta Time Travel: read as of a specific timestamp
df_delta_ts = spark.read \
    .format("delta") \
    .option("timestampAsOf", "2024-01-01") \
    .load("/output/transactions_delta")
print("Delta Time Travel (as of 2024-01-01):")
df_delta_ts.show()

# -------------------------------------------------------
# Write Modes
# -------------------------------------------------------

# overwrite: replaces all existing data
df.write.mode("overwrite").parquet("/output/transactions_parquet")

# append: adds new data to existing files
df.write.mode("append").parquet("/output/transactions_parquet")

# ignore: does nothing if data already exists
df.write.mode("ignore").parquet("/output/transactions_parquet")

# errorIfExists (default): throws error if data already exists
df.write.mode("error").parquet("/output/transactions_parquet")

# -------------------------------------------------------
# Control output file count
# -------------------------------------------------------

# Repartition to control number of output files
# 1 partition = 1 output file, useful for small datasets
df.repartition(1).write \
    .mode("overwrite") \
    .parquet("/output/transactions_single_file")

# coalesce reduces partitions without a full shuffle
# more efficient than repartition when reducing partition count
df.coalesce(2).write \
    .mode("overwrite") \
    .parquet("/output/transactions_two_files")

What to say: I always write production data in Parquet or Delta rather than CSV or JSON. Parquet is columnar which means when a query only needs three columns from a fifty column table, Spark reads only those three columns from disk rather than reading all fifty and discarding forty seven, which can reduce IO by ninety percent or more. Partitioning by a column like region or date on write allows partition pruning on read, where Spark skips entire directories of files based on the filter condition without even opening them.

Follow-up Questions & Answers

Q1. What is the difference between Parquet and Delta Lake? Why would you choose Delta over plain Parquet?

Parquet is a file format, a way of storing data on disk in a columnar compressed layout. It has no concept of transactions, versioning, or updates. Delta Lake is a storage layer built on top of Parquet files that adds a transaction log, which is a folder called _delta_log containing JSON files that record every operation performed on the table. This transaction log is what gives Delta its most important capabilities. ACID transactions mean multiple writers can write to the same table concurrently without corrupting each other's data. Time travel means you can query any previous version of the table by version number or timestamp. Schema enforcement means Delta rejects writes that do not match the table schema rather than silently accepting wrong data. Upserts using merge allow you to update existing rows and insert new ones in a single atomic operation, which plain Parquet cannot do at all. I choose Delta whenever the pipeline involves concurrent writes, needs audit history, requires upserts, or runs in an environment that supports it like Databricks, Azure Synapse, or open source Delta on any Hadoop compatible storage.

Q2. What is the difference between repartition and coalesce and when would you use each?

Both repartition and coalesce change the number of partitions in a DataFrame but they do so very differently.

# repartition(n): performs a full shuffle to redistribute data
# into exactly n partitions, evenly distributed
# use when increasing partitions or when you need even distribution
df_repartitioned = df.repartition(10)

# repartition by column: shuffles data so all rows with the
# same column value go to the same partition
# useful before a join or groupBy on that column
df_by_region = df.repartition(10, "region")

# coalesce(n): merges existing partitions without a full shuffle
# can only reduce partitions, cannot increase them
# more efficient than repartition when reducing
df_coalesced = df.coalesce(2)

What to say: repartition performs a full shuffle across the network which is expensive but produces perfectly balanced partitions and can both increase and decrease the partition count. coalesce avoids the full shuffle by merging existing partitions on the same executor where possible, making it much cheaper when you just want to reduce partitions, for example before writing a small result to a single output file. The rule I follow is use coalesce when reducing partition count, use repartition when increasing or when you need even distribution by a specific key.

Q3. What is partition pruning and how does partitioning on write enable it?

When you write a Parquet or Delta table partitioned by a column like region, Spark creates a separate subdirectory for each distinct value of that column. The directory structure looks like region equals North, region equals South, and so on. When a subsequent query filters on that column, Spark inspects the directory names rather than opening any files and immediately knows which subdirectories to skip entirely. This is partition pruning, the ability to eliminate entire partitions from consideration before reading a single byte of actual data.

# Write partitioned by region creates this folder structure:
# /output/transactions/region=North/part-00000.parquet
# /output/transactions/region=South/part-00000.parquet

df.write.partitionBy("region").parquet("/output/transactions")

# This query only reads the North folder, South is never touched
spark.read.parquet("/output/transactions") \
     .filter(F.col("region") == "North") \
     .show()

# Check if Spark is actually pruning by looking at the execution plan
spark.read.parquet("/output/transactions") \
     .filter(F.col("region") == "North") \
     .explain()
# Look for PartitionFilters in the plan output

What to say: Partition pruning is one of the most impactful performance optimisations available at zero code cost. Choosing the right partition column is critical though. You want a column with low to medium cardinality, like date, region, or department, not a high cardinality column like customer_id or transaction_id, since partitioning by a column with millions of distinct values creates millions of tiny directories which actually harms performance rather than helping it.

Q4. What are the different write modes in Spark and when would each be appropriate in a production pipeline?

# overwrite: deletes all existing data and writes fresh
# appropriate for a full reload pipeline where source is always complete
df.write.mode("overwrite").parquet("/output/path")

# append: adds new data alongside existing data without touching it
# appropriate for incremental pipelines loading new partitions
df.write.mode("append").parquet("/output/path")

# ignore: skips the write entirely if data already exists
# appropriate for idempotent pipelines where re-running should be safe
df.write.mode("ignore").parquet("/output/path")

# error (default): throws an exception if any data already exists
# appropriate as a safety guard during development
df.write.mode("error").parquet("/output/path")

What to say: In production pipelines I almost always use overwrite for full reloads and append for incremental loads. The danger with append is that re-running a failed job without cleaning up the partial output from the previous run leads to duplicate data, which is why Delta Lake's merge operation is so valuable in upsert scenarios since it handles the re-run case atomically. For pure append pipelines I make sure the pipeline is truly idempotent, meaning re-running produces the same result as running once, either by deduplicating after append or by using Delta's merge to handle duplicates.

Q5. How would you read only the latest partition from a date partitioned Parquet table without reading older data?

A very common production pattern where you only need today's data rather than the full history.

from datetime import date

# Get today's date as a string
today = str(date.today())
print(f"Reading partition for: {today}")

# Option 1: Read specific partition directory directly
df_today = spark.read.parquet(f"/output/transactions/txn_date={today}")

# Option 2: Read full table and filter, letting partition pruning handle it
df_today = spark.read \
    .parquet("/output/transactions") \
    .filter(F.col("txn_date") == today)

# Option 3: Find the latest available partition dynamically
import subprocess

# List partition directories and find the max date
partitions = spark.read.parquet("/output/transactions") \
                  .select("txn_date") \
                  .distinct() \
                  .orderBy(F.col("txn_date").desc())

latest_date = partitions.first()["txn_date"]
print(f"Latest available partition: {latest_date}")

df_latest = spark.read \
    .parquet("/output/transactions") \
    .filter(F.col("txn_date") == latest_date)

df_latest.show()

What to say: Option 1 is the most efficient since it reads exactly one partition directory with no scanning overhead at all. Option 2 is cleaner code that relies on Spark's partition pruning to achieve the same result. Option 3 is the most robust for a production pipeline since it dynamically discovers the latest available partition rather than hardcoding today's date, which protects against missing partitions if the upstream pipeline failed to load data for a particular day. I would use Option 3 in any pipeline that needs to gracefully handle upstream data delivery delays.

The Interviewer Says

Sometimes the built in PySpark functions are not enough and you need to apply custom Python logic to a DataFrame column. Walk me through what a UDF is, how to create one, what the performance implications are, and how a Pandas UDF improves on a regular UDF. I want to see you write both types.

How to Handle This in the Interview

Before writing anything, say: UDF stands for User Defined Function. It lets me apply arbitrary Python logic to a Spark DataFrame column. However regular UDFs come with a significant performance cost that I always flag in an interview before writing one, and I would explain why Pandas UDFs are almost always the better choice.

Step 1: Explain why regular UDFs are expensive

• Say: A regular Python UDF breaks Spark's internal Java and Scala execution engine. For each row, Spark has to serialise the data from JVM memory into Python objects, pass them to the Python process, run the function, serialise the result back to JVM memory, and continue. This serialisation round trip happens for every single row and is the main performance bottleneck.

Step 2: Explain how Pandas UDFs solve this

• Pandas UDFs use Apache Arrow for serialisation which is a columnar in-memory format that transfers data between JVM and Python in batches rather than row by row. This eliminates the per-row serialisation overhead and makes Pandas UDFs typically ten to one hundred times faster than regular UDFs.

Step 3: Know when to still use a regular UDF

• A regular UDF is acceptable when the DataFrame is small, when the logic cannot be expressed as a Pandas operation, or when the team is not familiar with Pandas and simplicity matters more than performance.

Step 4: Always try built-in functions first

• Say: Before writing any UDF I always check whether a built-in PySpark function can do the job, since built-in functions run entirely in the JVM with no Python serialisation overhead at all. UDFs should be the last resort, not the first instinct.

Solution

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import StringType, IntegerType, DoubleType
import pandas as pd

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

data = [
    ("T01", "C01", "laptop computer", 75000),
    ("T02", "C02", "mobile phone", 25000),
    ("T03", "C03", "tablet device", 30000),
    ("T04", "C04", "computer monitor", 15000),
    ("T05", "C05", "wireless keyboard", 5000),
]

df = spark.createDataFrame(
    data,
    ["txn_id", "customer_id", "product", "amount"]
)

# -------------------------------------------------------
# PART 1: Regular Python UDF
# -------------------------------------------------------

# Define a Python function
def capitalise_words(text):
    if text is None:
        return None
    return " ".join(word.capitalize() for word in text.split())

# Register as a UDF with return type
capitalise_udf = F.udf(capitalise_words, StringType())

# Apply UDF to a column
df_with_clean_product = df.withColumn(
    "product_clean",
    capitalise_udf(F.col("product"))
)
print("After applying regular UDF:")
df_with_clean_product.show(truncate=False)

# Register UDF for use in Spark SQL as well
spark.udf.register("capitalise_words_sql", capitalise_words, StringType())
df.createOrReplaceTempView("transactions")
spark.sql("SELECT txn_id, capitalise_words_sql(product) FROM transactions").show()

# -------------------------------------------------------
# PART 2: UDF using decorator syntax
# -------------------------------------------------------

@F.udf(returnType=StringType())
def categorise_amount(amount):
    if amount is None:
        return "Unknown"
    if amount >= 50000:
        return "High Value"
    elif amount >= 20000:
        return "Mid Value"
    else:
        return "Low Value"

df_with_category = df.withColumn(
    "amount_category",
    categorise_amount(F.col("amount"))
)
print("After applying decorator UDF:")
df_with_category.show()

# -------------------------------------------------------
# PART 3: Pandas UDF (Series to Series)
# -------------------------------------------------------
from pyspark.sql.functions import pandas_udf

# Pandas UDF processes an entire column batch as a Pandas Series
# Much faster than regular UDF due to Apache Arrow serialisation
@pandas_udf(StringType())
def capitalise_words_pandas(series: pd.Series) -> pd.Series:
    return series.str.title()

df_with_pandas_udf = df.withColumn(
    "product_clean_pandas",
    capitalise_words_pandas(F.col("product"))
)
print("After applying Pandas UDF:")
df_with_pandas_udf.show(truncate=False)

# -------------------------------------------------------
# PART 4: Pandas UDF for numerical transformation
# -------------------------------------------------------

@pandas_udf(DoubleType())
def apply_tax(amount: pd.Series) -> pd.Series:
    return amount * 1.18  # Apply 18% GST

df_with_tax = df.withColumn(
    "amount_with_tax",
    apply_tax(F.col("amount").cast(DoubleType()))
)
print("After applying tax Pandas UDF:")
df_with_tax.show()

# -------------------------------------------------------
# PART 5: Pandas UDF GroupedMap (applyInPandas)
# -------------------------------------------------------

# applyInPandas applies a Pandas function to each group
# useful for complex group-wise operations not possible with groupBy agg

def normalise_amount(pdf: pd.DataFrame) -> pd.DataFrame:
    pdf["normalised_amount"] = (
        (pdf["amount"] - pdf["amount"].min()) /
        (pdf["amount"].max() - pdf["amount"].min())
    ).round(4)
    return pdf

from pyspark.sql.types import StructType, StructField

result_schema = StructType(
    df.schema.fields + [
        StructField("normalised_amount", DoubleType(), True)
    ]
)

df_normalised = df.groupBy("customer_id").applyInPandas(
    normalise_amount,
    schema=result_schema
)
print("After applyInPandas normalisation:")
df_normalised.show()

What to say: I always try to solve the problem with built-in PySpark functions first since they run entirely in the JVM with no Python overhead. If I genuinely need custom Python logic I reach for a Pandas UDF over a regular UDF because Pandas UDFs use Apache Arrow to batch the data transfer between JVM and Python, eliminating the row-by-row serialisation cost of regular UDFs. The applyInPandas pattern is useful for complex group-wise operations like normalisation or custom aggregations that cannot be expressed with the standard groupBy agg functions.

Follow-up Questions & Answers

Q1. What are the main performance problems with regular Python UDFs in PySpark?

A regular Python UDF breaks out of Spark's native JVM execution engine for every single row. Spark serialises each row's data from JVM internal format into Python pickle format, sends it across a socket to a Python subprocess, runs the function, serialises the result back from Python pickle format to JVM format, and then continues with the next row. This serialisation and inter-process communication overhead happens once per row. On a DataFrame with one hundred million rows that is one hundred million serialise-call-deserialise cycles. Additionally regular UDFs are opaque to Spark's Catalyst optimiser, meaning the optimiser cannot look inside the function to apply any optimisations, which means predicate pushdown and other plan optimisations stop working at the UDF boundary.

Q2. How does Apache Arrow make Pandas UDFs faster?

Apache Arrow is an in-memory columnar data format that both the JVM and Python can read directly without any conversion. When Spark passes data to a Pandas UDF it serialises an entire batch of rows into Arrow columnar format once, passes the Arrow buffer to Python which reads it directly as a Pandas DataFrame or Series without any further conversion, runs the Pandas function on the whole batch, then passes the result Arrow buffer back to JVM. The key difference from a regular UDF is that the serialisation and transfer happens once per batch of thousands of rows rather than once per individual row. Batch sizes are configurable through spark.sql.execution.arrow.maxRecordsPerBatch which defaults to ten thousand rows per batch.

# Configure Arrow batch size
spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "50000")

# Enable Arrow for Pandas UDFs (usually enabled by default in modern Spark)
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

Q3. When would you use applyInPandas versus a regular groupBy agg?

groupBy agg is the right choice when the aggregation you need is expressible as a combination of standard functions like sum, avg, count, min, max, collect_list, and so on. applyInPandas is the right choice when the group-level computation requires custom logic that cannot be expressed with standard aggregation functions.

# applyInPandas use cases that groupBy cannot handle

# 1. Fitting a model per group and returning predictions
def fit_and_predict(pdf):
    from sklearn.linear_model import LinearRegression
    model = LinearRegression()
    model.fit(pdf[["feature"]], pdf["target"])
    pdf["prediction"] = model.predict(pdf[["feature"]])
    return pdf

# 2. Computing rolling statistics with a custom window logic
def custom_rolling(pdf):
    pdf = pdf.sort_values("date")
    pdf["custom_7day_avg"] = pdf["amount"].rolling(7, min_periods=1).mean()
    return pdf

# 3. Percentile normalisation within each group
def percentile_rank(pdf):
    pdf["pct_rank"] = pdf["amount"].rank(pct=True)
    return pdf

What to say: The practical rule is that if I find myself writing a Python function that takes a Pandas DataFrame as input and returns a Pandas DataFrame as output with group-level logic that is more than two or three lines, applyInPandas is probably the right tool. If the logic is a simple element-wise transformation on a single column, a Pandas UDF Series to Series is better. If it can be expressed with built-in functions, avoid UDFs entirely.

Q4. How would you test a UDF before deploying it in a production pipeline?

Testing UDFs in isolation before applying them to a full Spark DataFrame is important because UDF errors surface at execution time and can be hard to debug once they are running across hundreds of partitions.

# Test the underlying Python function directly first
# without involving Spark at all
def capitalise_words(text):
    if text is None:
        return None
    return " ".join(word.capitalize() for word in text.split())

assert capitalise_words("hello world") == "Hello World"
assert capitalise_words(None) is None
assert capitalise_words("") == ""
assert capitalise_words("UPPER CASE") == "Upper Case"
print("All unit tests passed")

# Then test on a small sample DataFrame before running on full data
sample_df = spark.createDataFrame(
    [("hello world",), (None,), ("",), ("UPPER CASE",)],
    ["product"]
)
capitalise_udf = F.udf(capitalise_words, StringType())
sample_df.withColumn("cleaned", capitalise_udf(F.col("product"))).show()

What to say: I always test the raw Python function with unit tests before registering it as a UDF, since a pure Python function is much easier to test and debug than a UDF running inside Spark. Once the Python function passes all edge cases including None inputs, empty strings, and unexpected data types, I register it as a UDF and test on a small sample DataFrame before applying it to the full production dataset. This catches any type mismatch or serialisation issues that only appear when the function runs inside Spark's execution engine.

Q5. Can you register a UDF and use it in Spark SQL? How?

Yes, UDFs can be registered in the Spark session's function registry and then called by name in any Spark SQL query, which is useful when the team uses SQL heavily alongside Python code.

# Register with a SQL-friendly name
spark.udf.register("capitalise", capitalise_words, StringType())
spark.udf.register("categorise_amount", lambda x: "High" if x and x >= 50000 else "Low", StringType())

# Create a temp view to query
df.createOrReplaceTempView("transactions")

# Use registered UDFs in Spark SQL
result = spark.sql("""
    SELECT
        txn_id,
        capitalise(product) AS product_clean,
        categorise_amount(amount) AS amount_category,
        amount
    FROM transactions
    WHERE amount > 10000
""")
result.show()

# Pandas UDFs can also be registered for SQL use
@pandas_udf(StringType())
def title_case(s: pd.Series) -> pd.Series:
    return s.str.title()

spark.udf.register("title_case", title_case)
spark.sql("SELECT title_case(product) FROM transactions").show()

What to say: Registering UDFs for SQL use is particularly valuable in organisations where data analysts write Spark SQL but the custom logic was written by a data engineer in Python. The engineer registers the UDF once at session startup and analysts can call it by name in any SQL query without needing to know it is a Python function underneath. In a Databricks or Hive metastore environment you can also persist UDF registrations so they survive across sessions, though for production pipelines I prefer re-registering them explicitly in the job initialisation code to avoid depending on state that may or may not exist in the metastore.