04.10 Advanced PySpark

Advanced PySpark is where DataFrame code becomes a distributed query program. The objective is not to memorize more functions; it is to understand how expressions become optimized plans and how those plans consume CPU, memory, network and storage.

1. Advanced PySpark mental model

Python DataFrame API
        |
        v
Column expressions / SQL expressions
        |
        v
Logical plan
        |
        v
Catalyst optimization
        |
        v
Physical plan
        +---------------------------+
        |                           |
        v                           v
Native Spark execution       Python boundary
        |                           |
        v                           v
JVM / Photon / executors     Python workers / Arrow
        |
        v
Shuffle / storage / output

Architect rule: prefer native Spark expressions whenever Spark can represent the business logic directly. This keeps the operation visible to the optimizer and usually avoids unnecessary Python serialization and execution boundaries.

2. Catalyst-aware PySpark

PySpark is a Python API, but a DataFrame transformation does not normally loop through rows in Python. The Python program builds an expression tree; Spark analyzes, optimizes and executes that plan across the cluster.

from pyspark.sql import functions as F

df = (
    orders
    .filter(F.col("status") == "COMPLETE")
    .select("customer_id", "amount")
    .groupBy("customer_id")
    .agg(F.sum("amount").alias("revenue"))
)

Think in this sequence:

Python API -> expressions -> logical plan -> optimized plan -> physical plan -> tasks

That is why two pieces of Python that look similar can have very different runtime behavior.

3. Explain plans before optimization

Use the plan as an engineering diagnostic:

df.explain()
df.explain("formatted")
df.explain("extended")
df.explain("cost")

Mode	Main use
simple	Quick physical-plan inspection
formatted	Readable physical-plan structure
extended	Parsed, analyzed, optimized and physical plans
cost	Plan plus available optimizer statistics

When investigating a slow transformation, look for Exchange, Sort, repeated scans, late filters, wide projections, expensive joins, global ordering and Python execution nodes.

Do not optimize the Python source first. Inspect the physical plan first.

4. Native column expressions

Most business transformations can remain inside Spark:

from pyspark.sql import functions as F

result = (
    df
    .withColumn(
        "net_amount",
        F.coalesce(F.col("gross_amount"), F.lit(0)) -
        F.coalesce(F.col("discount"), F.lit(0))
    )
    .withColumn(
        "status_group",
        F.when(F.col("status") == "C", "COMPLETE")
         .when(F.col("status") == "R", "REJECTED")
         .otherwise("OTHER")
    )
)

Useful native building blocks include conditional expressions, casts, string functions, date functions, JSON functions, array functions and aggregations. The more logic remains in the expression tree, the more execution choices Spark can make.

5. Complex and nested data

Modern lakehouse data often contains Struct, Array and Map values:

customer
  +-- profile
  |     +-- country
  |     +-- segment
  +-- orders[]
        +-- order_id
        +-- items[]
              +-- sku
              +-- quantity
              +-- price

Access nested fields directly:

df.select("customer.profile.country", "customer.profile.segment")

Do not flatten every nested structure automatically. Preserve hierarchy when it represents a meaningful business object and flatten at an analytical boundary.

6. Higher-order functions

Higher-order functions let Spark process arrays and maps without converting each element into a Python loop.

transform

df.select(F.transform("items", lambda x: x * 1.18).alias("items_with_tax"))

filter

df.select(F.filter("items", lambda x: x > 100).alias("high_value_items"))

exists

df.select(F.exists("items", lambda x: x > 1000).alias("has_large_item"))

aggregate

df.select(F.aggregate("items", F.lit(0), lambda acc, x: acc + x).alias("item_total"))

Use these when the business operation can stay within the existing row grain. This can be much cleaner than exploding an array and aggregating it back.

7. explode, posexplode and inline

Use generators when the business grain really needs to change.

df.select("order_id", F.explode("items").alias("item"))

Use posexplode when array position matters:

df.select("order_id", F.posexplode("items").alias("position", "item"))

For arrays of structs, inline can project the struct fields into rows:

df.select("order_id", F.inline("items"))

Architect warning: explode is a cardinality operation. One order with 100 items can become 100 rows. State the grain before and after every explosion.

8. JSON: parse once, use typed structures

Prefer this architecture:

Raw JSON string
      |
      v
from_json + explicit schema
      |
      v
Struct / Array / Map
      |
      v
Native expressions
      |
      v
Curated columns

from pyspark.sql.types import StructType, StructField, StringType, DoubleType

payload_schema = StructType([
    StructField("currency", StringType(), True),
    StructField("amount", DoubleType(), True),
])

parsed = df.withColumn(
    "payload_struct",
    F.from_json("payload", payload_schema)
)

result = parsed.select(
    "id",
    "payload_struct.currency",
    "payload_struct.amount"
)

Use explicit schemas when the JSON is a production contract. Design a quarantine path for malformed records instead of silently losing data.

9. selectExpr

selectExpr is useful when SQL-style expressions make a transformation clearer:

result = df.selectExpr(
    "customer_id",
    "amount * 1.18 AS amount_with_tax",
    "CASE WHEN status = 'C' THEN 'COMPLETE' ELSE 'OTHER' END AS status_group"
)

Use it deliberately. In reusable Python libraries, explicit functions expressions can be easier to refactor and test.

10. Reusable transformation design

A production PySpark codebase should not be one giant notebook. Prefer composable transformations:

def standardize_orders(df):
    return (
        df
        .withColumnRenamed("cust_id", "customer_id")
        .withColumn("amount", F.col("amount").cast("decimal(18,2)"))
    )

def validate_orders(df):
    return df.filter(
        F.col("customer_id").isNotNull() &
        F.col("amount").isNotNull()
    )

def enrich_orders(df, customers):
    return df.join(F.broadcast(customers), "customer_id", "left")

Benefits include unit testing, code review, reuse, deployment, observability and easier debugging.

11. Cache and persist

Cache only when a computed DataFrame is reused and the measured savings justify memory/storage cost.

from pyspark import StorageLevel

prepared = expensive_transformation(df)
prepared.persist(StorageLevel.MEMORY_AND_DISK)
prepared.count()  # materialize
report_a = prepared.groupBy("region").count()
report_b = prepared.groupBy("product").count()
prepared.unpersist()

Ask: Is it reused? Is recomputation expensive? Does it fit? Will it be reused before eviction? Did measured runtime improve?

12. Checkpointing vs caching

Technique	Primary purpose
Cache / persist	Reuse computed data efficiently
Checkpoint	Cut lineage and create a stable recovery boundary

Checkpointing introduces storage I/O; it is not a free performance switch. Use it for long or iterative lineage when a stable boundary is useful.

13. Partitioning

Partition count is a performance control, not a magic number.

df.rdd.getNumPartitions()
df.repartition(200, "customer_id")
df.coalesce(20)

Operation	Typical effect
repartition	Redistributes data and normally causes a shuffle
coalesce	Reduces partitions with less redistribution

Do not add repartition calls just in case. Every unnecessary shuffle can add network, CPU and disk work.

14. Avoid accidental driver work

These operations can move a large result to the driver:

df.collect()
df.toPandas()
df.take(1000000)

Use them only for deliberately bounded results. For inspection, prefer:

df.limit(20).show(truncate=False)

Architect rule: keep large data distributed; bring only controlled results to the driver.

15. Advanced performance pattern

Consider a pipeline that filters orders, projects columns, joins a customer dimension and aggregates revenue. Optimize in this order:

1. Establish row grain and cardinality.
2. Filter early where it is semantically safe.
3. Project only required columns.
4. Validate join cardinality and strategy.
5. Look for unnecessary Exchange operators.
6. Check skew and spill in the Spark UI.
7. Evaluate AQE/runtime statistics.
8. Measure before and after each change.

Use the loop:

Measure -> inspect plan -> identify expensive boundary -> change one thing -> measure again

AQE can use runtime statistics to coalesce post-shuffle partitions and address some skew scenarios. Avoid replacing evidence with arbitrary Spark configuration changes.

16. Production Bronze -> Silver -> Gold

BRONZE
Raw / replayable
    |
    | parsing + schema validation + technical metadata
    v
SILVER
Typed / clean / conformed
    |
    | business transformations + joins + quality
    v
GOLD
Business-ready
    +--> BI
    +--> ML
    +--> APIs
    +--> downstream products

Advanced PySpark is especially important in Silver and Gold, where correctness, grain, performance and maintainability must all be explicit.

17. Enterprise SAP example

Suppose a raw sales payload contains an order, customer structure and an items array.

{
  "sales_order": "50001234",
  "customer": {"id": "0000123456", "country": "IN"},
  "items": [
    {"material": "MAT01", "qty": 10, "price": 125.50},
    {"material": "MAT02", "qty": 2, "price": 950.00}
  ]
}

Production flow:

Raw event -> from_json(schema) -> validate -> explode(items) -> item Silver -> enrich -> Gold

Keep identifiers such as SAP customer numbers as strings when leading zeros are meaningful. After explode, the grain is sales-order-item, not sales-order.

18. Common anti-patterns

Python row loops

Never collect a large dataset just to iterate in Python. Use distributed DataFrame operations.

UDF for simple logic

Use native expressions for trim, upper, arithmetic, conditional logic and similar operations.

Explode everything

Explosion can multiply rows and make downstream operations much more expensive.

Cache everything

Cache is a workload optimization, not a default.

Repartition everywhere

Repartition is often a shuffle; use it only when the new distribution solves a real problem.

Repeated JSON parsing

Parse once into a typed structure and reuse it.

Blind collect / toPandas

Bound the result before crossing to the driver.

19. Hands-on Lab 1 — Native vs UDF

Create a DataFrame containing customer, amount, discount and status. Implement the same transformation with native expressions and a Python UDF.

Run:

native.explain("formatted")
udf_version.explain("formatted")

Compare physical plans, execution time and Python execution boundaries. The goal is to recognize when a UDF changes execution architecture.

20. Hands-on Lab 2 — Nested JSON

Build a JSON payload with customer, order and items. Then:

1. Define an explicit schema.
2. Parse with from_json.
3. Access nested fields.
4. Apply transform to an array.
5. Explode only when item grain is required.
6. Aggregate back to order grain.
7. Write the result to Delta.

Document the grain at every step.

21. Hands-on Lab 3 — Plan-driven optimization

Start with a deliberately inefficient pipeline containing select(*), an arbitrary repartition, late filtering, a join and cache. Then optimize it using evidence.

Change	Why	Plan impact
Filter early	Reduce rows	Smaller input
Select columns	Reduce width	Narrower plan
Remove unnecessary repartition	Remove shuffle	Fewer Exchanges
Remove unnecessary cache	Reduce memory pressure	Simpler execution

The objective is evidence-driven optimization, not a magic configuration.

22. Troubleshooting workflow

Step 1 — Establish grain

What does one row represent?

Step 2 — Check cardinality

Did a join or explode multiply the data?

Step 3 — Inspect the physical plan

Look for Exchange, Sort, Broadcast, Python execution and repeated scans.

Step 4 — Check Spark UI

Look for long stages, skewed tasks, spill, shuffle read/write and executor memory pressure.

Step 5 — Test one hypothesis

Do not change ten settings simultaneously.

Step 6 — Re-measure

Optimization without measurement is guesswork.

23. Architect decision matrix

Requirement	Preferred approach
Simple column logic	Native Spark expression
Array transformation	Higher-order function
One row per array element	explode
Nested JSON	from_json + explicit schema
Reuse expensive DataFrame	Cache/persist after measurement
Cut long lineage	Checkpoint when justified
Reduce output partitions	coalesce when appropriate
Redistribute by key	repartition
Small dimension join	Consider broadcast
Large distributed join	Let Spark/AQE choose, then inspect
Driver-side result	Only if explicitly bounded

24. Interview questions

Why can PySpark DataFrame code be distributed even though the API is Python? Because DataFrame operations construct Spark expressions and execution plans; Spark executes the plan across the cluster.

Why prefer native expressions over Python UDFs? Native expressions remain visible to Spark's optimizer and avoid an unnecessary Python execution boundary.

Cache vs checkpoint? Cache/persist targets reuse and speed; checkpoint creates a stable lineage boundary and may involve storage I/O.

When should you use explode? When downstream processing genuinely requires a new row for each array element.

How would you optimize a slow PySpark pipeline? Establish grain, inspect the physical plan, reduce rows and columns, validate joins, eliminate unnecessary shuffles, use runtime statistics/AQE, benchmark each change and verify the Spark UI.

25. Production cheat sheet

ADVANCED PYSPARK
|
+-- Think in plans, not Python statements
+-- Prefer native expressions
+-- Keep large data distributed
+-- Inspect explain("formatted")
+-- Track row grain after major operations
+-- Treat explode as a cardinality change
+-- Parse JSON once with an explicit contract
+-- Use higher-order functions for arrays/maps
+-- Cache only for measured reuse
+-- Checkpoint only when lineage requires it
+-- Repartition intentionally
+-- Avoid accidental driver collection
+-- Measure before and after optimization
+-- Use Spark UI to validate hypotheses

Official References

• Apache Spark SQL Performance Tuning
• Apache Spark SQL and DataFrame Guide
• Apache Spark Tuning Guide
• Databricks PySpark Functions
• Databricks from_json

Key Takeaway

Advanced PySpark is the discipline of connecting business requirement -> data grain -> expressions -> logical plan -> optimized physical plan -> shuffle/CPU/memory/network -> cluster execution -> measured outcome.

Once you can reason through that chain, PySpark becomes an engineering discipline rather than a collection of APIs.

Next: 04.11 PySpark Performance Engineering.