backend / data engineering / spark databricks / 02_pyspark_gotchas_and_databricks.md

PySpark Gotchas and Databricks

7 interview angles 7 min read source

PySpark Gotchas and Databricks

Practical PySpark for Python backends and how Databricks fits.

The PySpark gap

PySpark is Spark’s Python binding. Most things work like Scala Spark. A few don’t, or have quirks.

Python ↔ JVM cost

[JVM executor]  ←(serialize)→ [Python worker]

Every Python UDF / row-level Python operation crosses this boundary. Costly. For aggregations / joins / filters / transformations, use Spark’s built-in functions:

# Slow — Python UDF
from pyspark.sql.functions import udf
@udf("integer")
def double_it(x):
    return x * 2
df.withColumn("doubled", double_it("amount"))

# Fast — Spark expression
df.withColumn("doubled", df["amount"] * 2)

10-100× difference for the operation.

pandas UDFs (vectorized)

When you must use Python, batch rows via pandas UDFs:

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def normalize(values: pd.Series) -> pd.Series:
    return (values - values.mean()) / values.std()

df.withColumn("z_score", normalize("amount"))

Behind the scenes: Arrow batches passed between JVM and Python. Much less serialization overhead.

groupedmap pandas UDFs let you process each group as a pandas DataFrame:

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def process_group(pdf):
    pdf["normalized"] = (pdf["value"] - pdf["value"].mean()) / pdf["value"].std()
    return pdf

df.groupBy("user_id").apply(process_group)

For per-group custom logic where Spark’s built-ins don’t fit.

Common PySpark patterns

Read-write parquet on S3

df = spark.read.parquet("s3://bucket/orders/")
df.write.mode("overwrite").parquet("s3://bucket/output/")

Spark partitions by S3 prefix. To partition output:

df.write.partitionBy("year", "month").parquet("s3://bucket/orders/")
# Files at: s3://bucket/orders/year=2026/month=05/...

Subsequent reads with predicates can prune partitions (spark.read.parquet(...).where("year = 2026") reads only year=2026 files).

Reading from JDBC databases

df = (
    spark.read
    .format("jdbc")
    .option("url", "jdbc:postgresql://host:5432/db")
    .option("dbtable", "orders")
    .option("user", "...")
    .option("password", "...")
    .option("numPartitions", 10)
    .option("partitionColumn", "id")
    .option("lowerBound", 1)
    .option("upperBound", 1000000)
    .load()
)

numPartitions + partitionColumn enables parallel reads (10 workers fetch different ID ranges). Without partitioning, one executor pulls the whole table single-threaded.

Writing to JDBC

df.write.format("jdbc").mode("append").option(...).save()

Many parallel writers can saturate the DB. Use numPartitions to limit, or write to S3 / Parquet and load to DB via the DB’s bulk-load tool (Postgres COPY, Redshift COPY).

Caching

df = expensive_computation()
df.cache()                  # persist in memory (and spill to disk if needed)
df.count()                  # action triggers caching
df.write.parquet(...)       # uses cached data
df.groupBy(...).count()     # also uses cached data

df.unpersist()              # free

Cache only when the same DataFrame is materialized multiple times. Caching once-used data wastes memory.

Storage levels:

  • MEMORY_ONLY — fast, evicts on memory pressure.
  • MEMORY_AND_DISK (default for DataFrame .cache()) — spills to disk.
  • DISK_ONLY — slow but always available.

SparkSession in production

spark = (
    SparkSession.builder
    .appName("orders-etl")
    .config("spark.sql.shuffle.partitions", 200)
    .config("spark.sql.adaptive.enabled", "true")
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
    .getOrCreate()
)

Key configs:

  • spark.sql.shuffle.partitions (default 200) — partitions after shuffles. Tune for your data size.
  • spark.sql.adaptive.enabled — AQE; major win for real workloads.
  • spark.executor.memory / cores — sizing.
  • spark.sql.autoBroadcastJoinThreshold — broadcast threshold.

When .collect() and .toPandas() bite

df.collect()      # brings the whole DataFrame to the driver
df.toPandas()     # same, then converts to pandas

Both load all data into the driver’s memory. For a 100 GB DataFrame, you OOM the driver.

Use only for small results:

  • After groupBy().count() returning N rows where N is small.
  • After .limit(N) with explicit small N.
  • For previews via .show() (which only fetches first 20 rows).

For larger results, write to S3 and process downstream.

Streaming basics

Spark Structured Streaming — micro-batch (and continuous-mode beta):

df = (
    spark.readStream
    .format("kafka")
    .option("kafka.bootstrap.servers", "...")
    .option("subscribe", "orders")
    .load()
)

result = df.selectExpr("CAST(value AS STRING) as raw") \
           .selectExpr("from_json(raw, 'order_id STRING, amount DOUBLE') as o") \
           .select("o.*")

query = (
    result.writeStream
    .format("parquet")
    .option("path", "s3://bucket/orders/")
    .option("checkpointLocation", "s3://bucket/checkpoints/orders/")
    .trigger(processingTime="1 minute")
    .start()
)
query.awaitTermination()

Reads Kafka in micro-batches (every 1 minute here), processes, writes Parquet. Checkpoints state for resume.

Properties:

  • Exactly-once via checkpoints (writes are idempotent or transactional, depending on sink).
  • Watermarks for late data: df.withWatermark("event_time", "10 minutes").
  • Windowing: window("event_time", "5 minutes").

For low-latency streaming (sub-second), Flink or Kafka Streams is better. Spark Structured Streaming is “small batch interval” — typical ~1-10 second latency.

Databricks

Managed Spark platform from the team behind Spark. Adds:

Unity Catalog

Three-layer hierarchy: catalog → schema → table. Centralized metadata, access control, lineage.

SELECT * FROM main.sales.orders;
GRANT SELECT ON main.sales.orders TO `data-analyst-group`;

Replaces the older Hive Metastore. Works across Databricks workspaces; integrates with Delta Sharing for cross-org data exchange.

Jobs and Workflows

# Define jobs in JSON / Terraform / Asset Bundles
# Trigger on cron / event / file arrival
# Multi-task workflows with dependencies

Replaces Airflow / Step Functions for Databricks-native pipelines. Cheaper if all your compute is on Databricks.

Delta Live Tables (DLT)

Declarative pipeline framework on top of Delta Lake. You declare tables; DLT manages dependencies, refresh, schema evolution, quality checks.

import dlt

@dlt.table
def bronze_orders():
    return spark.readStream.format("kafka")...

@dlt.table
@dlt.expect("valid_amount", "amount > 0")
def silver_orders():
    return dlt.read_stream("bronze_orders").filter(...)

@dlt.table
def gold_orders_per_user():
    return dlt.read("silver_orders").groupBy("user_id").agg(...)

DLT handles the DAG, monitoring, retries, schema. Higher-level than raw Spark.

Databricks SQL warehouses

Photon-engine SQL endpoints — Spark queries served via SQL, optimized for BI tool latency. Trade-off vs raw Spark: faster for SQL queries, less flexible than full Spark.

Photon

Databricks’ C++ vectorized query engine; replaces parts of the JVM execution path. 2-3× faster for typical analytic queries. Default in newer runtimes.

Auto Loader

df = (
    spark.readStream
    .format("cloudFiles")
    .option("cloudFiles.format", "json")
    .load("s3://bucket/incoming/")
)

Watches an S3 prefix, processes new files incrementally. Replaces ad-hoc “list files since last run” code.

Databricks vs raw Spark on EMR/Kubernetes

Databricks Raw Spark (EMR / K8s)
Setup minutes hours-days
Ops burden minimal significant
Cost premium (~30% over compute) base AWS / GCP cost
Features Unity Catalog, DLT, Jobs, Photon what you build
Lock-in Databricks runtime, Delta Live Tables none
Notebooks first-class install Jupyter / Zeppelin
Performance Photon win on SQL base Spark

For teams without dedicated platform engineering: Databricks. For cost-sensitive at huge scale or those wanting portable open-source stacks: EMR / GKE Dataproc / self-managed Spark on K8s.

Common production issues

Skew

One key dominates → one task takes forever. Enable AQE skew handling; salt keys; multi-stage aggregation.

Small files problem

10,000 tiny files in your write output → next read pays huge per-file overhead. Use coalesce(N) or repartition(N) before write, where N targets ~100-200 MB per file.

For Delta: OPTIMIZE my_table ZORDER BY (...) compacts small files.

Driver OOM on .collect() or .toPandas()

Don’t pull big results to driver. Write to S3 instead.

Long-running shuffles

spark.sql.shuffle.partitions = 200 (default) might be too few for big data → each partition is huge → spills to disk. Or too many for small data → scheduling overhead.

Tune for ~100-200 MB per shuffle partition.

Memory pressure on executors

ExecutorLostFailure (executor lost due to memory)

Bump spark.executor.memory, reduce partitions per executor (spark.executor.cores), or restructure the job.

Reading too many partitions

spark.read.parquet("s3://bucket/") over a 100k-file dataset is slow to start (S3 listing). Push partition pruning into the read:

spark.read.parquet("s3://bucket/").filter("year = 2026 AND month = 5")

Spark applies the filter to partition discovery if the data is partitioned by year / month. Otherwise full listing.

When to use Spark vs alternatives

Data size Tool
< 10 GB Polars / DuckDB on one machine
10-100 GB single machine Polars (streaming) / DuckDB
100 GB - few TB Spark on Databricks / EMR
> few TB Spark, possibly with Iceberg / Delta
Real-time streaming Flink / Kafka Streams (sub-second), or Spark Structured Streaming (~seconds)
Interactive SQL on big data Trino / Athena / BigQuery / Snowflake

For Python backend roles, you’re more likely to consume Spark output (read Parquet) than write Spark jobs. Knowing the architecture helps you understand why pipelines are shaped the way they are.

Interview angle

  • “PySpark UDF performance?” — regular UDFs serialize per row JVM↔Python — slow. Use built-in Spark functions when possible; pandas UDFs (vectorized, Arrow-based) when you need Python; regular UDFs only as last resort.
  • “What’s the shuffle partition default and when do you tune it?”spark.sql.shuffle.partitions = 200 by default. Too few for big data: huge partitions spill to disk. Too many for small data: scheduling overhead. Target ~100-200 MB per shuffle partition.
  • “Why is .collect() dangerous?” — pulls all data to the driver’s memory. For a 100 GB DataFrame, driver OOMs. Only use for small results; for big output, write to S3.
  • “AQE — what does it do?” — Adaptive Query Execution (Spark 3.0+). Re-plans at runtime based on actual statistics: adjusts shuffle partition count, converts sort-merge to broadcast when one side turns out small, handles skew. Default-on in newer versions; major real-world win.
  • “What’s Databricks Photon?” — C++ vectorized query engine replacing parts of JVM execution. 2-3× faster for analytics queries. Default in newer Databricks runtimes.
  • “Delta Live Tables — what’s the value?” — declarative pipeline on top of Delta Lake. You define tables and dependencies; DLT manages DAG, retries, schema, quality expectations. Higher-level than raw Spark; opinionated; locks you to Databricks.
  • “Spark Structured Streaming latency?” — micro-batch with typical 1-10 second latency. For sub-second, prefer Flink or Kafka Streams. For “near real-time” with strong consistency and SQL, Structured Streaming is fine.