Spark DataFrame API, in depth: column expressions, null logic, joins, aggregations, windows and the plans they build

What a DataFrame is: a plan, not a container

Calls such as select, filter, join and groupBy are transformations. They return a new DataFrame immediately and are checked only for names and types. Calls such as count, collect, show and write are actions. They send the plan to the optimizer, which pushes filters into scans, prunes unused columns and picks join strategies, and then run it.

Transformations come in two kinds. Narrow ones, like filters and projections, work on each partition independently. Wide ones, like grouping and most joins, need rows with the same key on the same executor, which means a shuffle across the network. Shuffles are where most of the time and most of the failures in a Spark job live; Spark shuffle explains the mechanics. Use df.explain("formatted") to see which kind you wrote; reading EXPLAIN plans covers the output.

Columns are expressions

A Column is not a vector of values. It is an expression tree that will be evaluated per row later. F.col("price") * 1.2 builds a multiply node; it does not compute anything. This is why you combine conditions with &, | and ~ rather than Python's and and or, and why each condition needs parentheses: Python operator precedence applies to the tree you are building.

from pyspark.sql import functions as F

orders = spark.read.parquet("s3://shop/orders/")   # lazy: reads only the footer schema

enriched = (
    orders
    .filter((F.col("status") == "PAID") & (F.col("amount") > 0))
    .select(
        "order_id", "customer_id", "amount",
        F.to_date("created_at").alias("order_date"),
        F.when(F.col("amount") >= 100, "large")
         .when(F.col("amount") >= 20, "medium")
         .otherwise("small").alias("size_band"),
    )
    .withColumns({
        "amount_eur": F.round(F.col("amount") * F.lit(0.92), 2),
        "is_weekend": F.dayofweek("order_date").isin(1, 7),
    })
)

select is the general tool: it chooses, renames and computes columns in one projection. withColumn adds or replaces one column, and withColumns (Spark 3.3 and later) does several at once. Prefer built-in functions in pyspark.sql.functions over Python code: they run inside the JVM, are visible to the optimizer and benefit from whole-stage code generation.

Null logic: the rule that silently drops rows

Spark follows SQL three-valued logic. Any comparison with NULL yields NULL, not false, and filter keeps only rows where the condition is true. So df.filter(F.col("country") != "US") drops every row whose country is null, which is rarely what the author meant.

The same logic applies elsewhere. count("col") skips nulls while count("*") does not. when without otherwise yields null for unmatched rows. Sums and averages ignore nulls, so an average over a column that is half null is an average of the other half.

Spark 4.0 changed the defaults: spark.sql.ansi.enabled is now true, so an invalid cast, a division by zero or an integer overflow raises an error instead of quietly producing null. That is safer, but jobs upgraded from Spark 3 can start failing on dirty data that used to pass. Where bad input is expected, say so explicitly with F.expr("try_cast(raw AS INT)") and then count the nulls it produced.

Aggregation

daily = (
    enriched
    .groupBy("order_date", "size_band")
    .agg(
        F.count("*").alias("orders"),
        F.sum("amount_eur").alias("revenue_eur"),
        F.approx_count_distinct("customer_id").alias("customers_approx"),
    )
)

groupBy triggers a shuffle on the grouping keys, but Spark first does a partial aggregation on each partition, so only one row per key per partition crosses the network. Exact countDistinct cannot be partially reduced the same way and is much more expensive on large data; approx_count_distinct uses a HyperLogLog sketch with a configurable error. pivot is convenient for reports; pass the list of values explicitly, otherwise Spark runs an extra job to discover them.

Joins: types, keys and nulls

Every join answers two questions: which rows survive, and how the keys are matched.

Join on a column name, orders.join(customers, on="customer_id", how="left"), and Spark keeps one customer_id column. Join on an expression, orders.customer_id == customers.customer_id, and you get two columns with the same name, which later fail as ambiguous references. Null keys never match under ==; use eqNullSafe if null should match null.

Strategy matters as much as semantics. A small dimension table should be broadcast to every executor with F.broadcast(customers) so the large side is not shuffled. When one key holds a large share of rows, one task does most of the work; adaptive execution can split such partitions, as described in Spark AQE.

Window functions and the default frame

A window function computes a value per row from a group of related rows without collapsing them. The classic use is keeping the latest record per key.

from pyspark.sql.window import Window

w_latest = Window.partitionBy("customer_id").orderBy(F.col("updated_at").desc(), F.col("event_id").desc())
latest = (
    events.withColumn("rn", F.row_number().over(w_latest))
          .filter("rn = 1").drop("rn")
)

# Running total: be explicit about the frame
w_running = (Window.partitionBy("customer_id").orderBy("order_date")
             .rowsBetween(Window.unboundedPreceding, Window.currentRow))
running = enriched.withColumn("cum_spend", F.sum("amount_eur").over(w_running))

Two details prevent wrong answers. First, add a tie-breaker to the ordering; with ties, row_number picks an arbitrary winner and reruns can disagree. Second, when a window has orderBy but no explicit frame, the default frame is a range from the start of the partition to the current row, so rows with equal ordering values are summed together. Two orders on the same date both show the full day's total. Use rowsBetween when you mean row by row.

Worked example: from raw orders to a tested daily table

Suppose the raw table has 400 million orders, with 3% duplicated by an at-least-once ingester, 2% with a null country, and the analytics team wants daily revenue by country, excluding the US. A first version, filter(F.col("country") != "US") after a plain dropDuplicates(), returned a total only about 1% above the finance figure, so nobody noticed. It was two errors partly cancelling: the duplicates differed in an ingestion timestamp, so whole-row deduplication kept both copies and inflated revenue by about 3%, while the filter dropped the null-country rows and removed about 2%.

The fixed version deduplicates on the business key with the window above, maps null to an explicit "UNKNOWN" bucket with F.coalesce, and filters with eqNullSafe. A test pins the behaviour so it cannot regress:

from pyspark.testing import assertDataFrameEqual

def daily_revenue(orders):
    dedup = (orders.withColumn("rn", F.row_number().over(
                 Window.partitionBy("order_id").orderBy(F.col("ingested_at").desc())))
             .filter("rn = 1"))
    return (dedup
            .withColumn("country", F.coalesce("country", F.lit("UNKNOWN")))
            .filter(~F.col("country").eqNullSafe("US"))
            .groupBy("order_date", "country")
            .agg(F.sum("amount").alias("revenue")))

def test_daily_revenue(spark):
    rows = [("o1", "2026-10-01", None, 10.0, 1),
            ("o1", "2026-10-01", None, 10.0, 2),   # duplicate delivery
            ("o2", "2026-10-01", "US", 5.0, 1),
            ("o3", "2026-10-01", "DE", 7.0, 1)]
    df = spark.createDataFrame(rows, "order_id string, order_date string, country string, amount double, ingested_at int")
    expected = spark.createDataFrame(
        [("2026-10-01", "UNKNOWN", 10.0), ("2026-10-01", "DE", 7.0)],
        "order_date string, country string, revenue double")
    assertDataFrameEqual(daily_revenue(df), expected)   # row order ignored by default

assertDataFrameEqual arrived in PySpark 3.5; on older versions compare sorted collect() output. Small, explicit inputs with the awkward cases in them (a duplicate, a null, the excluded value) catch more bugs than large samples do.

Where the plan starts and ends: reads and writes

The read is the first node of every plan, and the optimizer works hardest there. With a columnar source such as Parquet, a select of three columns reads only those columns from storage, and a filter on a partition column skips whole directories. Both happen automatically as long as the filter is a plain expression Spark can understand; wrap the column in a Python UDF and the pushdown is lost, because the optimizer cannot see inside it.

For row formats such as CSV and JSON, pass the schema explicitly instead of letting Spark infer it. Inference costs an extra pass over the data and produces a different schema when a new file contains an unexpected value.

schema = "order_id string, customer_id string, amount double, status string, created_at timestamp"
raw = spark.read.schema(schema).json("s3://shop/raw/orders/2026-10-02/")

(daily.repartition("order_date")              # one shuffle, a few files per date
      .write.mode("overwrite")
      .partitionBy("order_date")
      .parquet("s3://shop/marts/daily_revenue/"))

On the write side, the number of output files equals the number of tasks writing to each partition directory, so a job with 2,000 tasks writing 30 dates can produce 60,000 small files. Repartitioning by the partition column first, as above, trades one shuffle for a sane file count. Be careful with mode("overwrite") on partitioned tables: by default it replaces the whole table, and spark.sql.sources.partitionOverwriteMode=dynamic is what limits it to the partitions present in the data being written.

Traps that cost hours

• withColumn in a loop. Each call adds a projection to the plan, and analysis cost grows with plan size. Adding hundreds of columns this way can stall the driver before any task runs. Build the list of expressions and use one select or withColumns.
• Python UDFs. A row-at-a-time Python UDF serialises every row to a Python worker and back and is opaque to the optimizer. Use built-ins first, then vectorised pandas UDFs, which move data in Arrow batches; Spark Arrow explains the transfer.
• collect and toPandas on large data. Both pull everything to the driver. Aggregate or limit first, or write the result out.
• Caching by reflex. cache() helps only when a DataFrame is reused by several actions, and costs executor memory. Unpersist when done.
• Very long lineage. Iterative algorithms that reassign a DataFrame thousands of times build enormous plans. Truncate with checkpoint() or localCheckpoint().
• Errors far from their cause. Because execution is lazy, a bad cast fails at the action. Run an action on a small sample after each stage while developing.

DataFrame, SQL or Dataset

DataFrame code and Spark SQL compile to the same plans and perform the same; choose by who maintains the code. DataFrames compose well in functions and tests, while SQL is easier for analysts to review. Mixing them is fine: register a view with createOrReplaceTempView and call spark.sql. Typed Datasets in Scala and Java add compile-time types at the cost of serialisation in lambdas; in Python, DataFrames are the only option.

What to do next

1. Run explain("formatted") on your most expensive job and mark each exchange; each is a shuffle you pay for.
2. Search your codebase for != and isin filters on nullable columns and decide explicitly what nulls should do.
3. Replace withColumn loops with a single select or withColumns.
4. List every Python UDF and replace it with built-in functions or a pandas UDF.
5. Add a tie-breaker to every row_number window and an explicit frame to every running aggregate.
6. Before moving to Spark 4, run your jobs with spark.sql.ansi.enabled=true on recent data and fix the casts that fail.
7. Write one assertDataFrameEqual test per transformation, with a duplicate, a null and an edge value in the input.