Back to Apache Spark: Core
Q39Question 39 of 72

Spark write-time partitioning (`df.write.partitionBy`) and partition pruning

Key Concepts
  • partitionBy writes data into Hive-style value-named directories so queries filtering on that column skip irrelevant directories entirely.
  • The partition column's value lives in the directory name, not inside the data files, saving storage.
  • Partition pruning at read time eliminates whole directories before a single row is read, dramatically reducing I/O.
  • Multi-column partitioning creates a nested directory tree; put the most-filtered column first.
  • Low-cardinality columns like branch, country, or month are ideal; high-cardinality columns cause directory explosion and small-file problems.
  • Combine partitioning with bucketing when you need both directory-level pruning and hash-based co-location.

Definition

`partitionBy` writes data into value-named directories so queries filtering on that column read only the matching directories and skip everything else.
  • Write-time partitioning – df.write.partitionBy("col") organizes output into a directory hierarchy keyed by column values, using Hive-style column=value directory names.
  • Partition pruning – when a later query filters on the partition column, Spark reads only the matching directories and skips the rest entirely, before any row is read.
  • Column placement – the partition column's value is encoded in the directory path and removed from the data files, saving per-row storage.
  • Cardinality rule – works best for low-cardinality columns; high cardinality causes directory explosion and the small-file problem.

Imagine the National Bank of India stores transaction records in a giant filing room. Without partitioning, every record for every branch is dumped into one enormous drawer, and finding Mumbai's records means reading everything. Write-time partitioning is like labelling one drawer per branch: 'Mumbai', 'Delhi', 'Chennai'. Now when the auditor asks only for Mumbai, the clerk opens just the Mumbai drawer and ignores the rest. That skipping of irrelevant drawers is partition pruning.

Understanding the Concepts

What is write-time partitioning?

Write-time partitioning instructs Spark to physically organize output files into a directory hierarchy keyed by the values of one or more columns. Each distinct value becomes a subdirectory. The partition column's value is encoded in the directory path and removed from the data files.

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

spark = SparkSession.builder.appName("NBI-PartitionBy").getOrCreate()

txn = spark.createDataFrame([
    ("SAV2001", "Mumbai",  50000.00, "2026-01-15"),
    ("SAV2002", "Delhi",   12000.00, "2026-01-15"),
    ("CUR3001", "Chennai", 300000.00, "2026-01-16"),
    ("SAV2001", "Mumbai",  75000.00, "2026-01-16")
], ["account_id", "branch", "amount", "txn_date"])

# Write partitioned by branch
txn.write.mode("overwrite").partitionBy("branch").parquet("/bank/txn-by-branch/")
TEXT
Output directory structure:
/bank/txn-by-branch/
  branch=Mumbai/
    part-00000.parquet
  branch=Delhi/
    part-00000.parquet
  branch=Chennai/
    part-00000.parquet

How does partition pruning work?

Partition pruning happens at read time. When you filter on the partition column, Catalyst recognizes the filter matches the directory structure. Instead of listing and reading every directory, it lists only the matching ones and skips the rest entirely. This is data skipping at the directory level: the irrelevant data is never read at all.

PYTHON
# Read the partitioned data
df = spark.read.parquet("/bank/txn-by-branch/")

# Filter on partition column: only branch=Mumbai directory is read
mumbai = df.filter(F.col("branch") == "Mumbai")
mumbai.explain()
mumbai.show()
TEXT
== Physical Plan ==
*(1) ColumnarToRow
+- FileScan parquet [account_id,amount,txn_date,branch]
     PartitionFilters: [branch = Mumbai]
     PushedFilters: []
     ReadSchema: struct<account_id:string,amount:double,txn_date:string>

+----------+--------+----------+------+
|account_id|  amount|  txn_date|branch|
+----------+--------+----------+------+
|   SAV2001| 50000.0|2026-01-15|Mumbai|
|   SAV2001| 75000.0|2026-01-16|Mumbai|
+----------+--------+----------+------+

How does multi-column partitioning work?

You can partition by multiple columns, creating a nested directory tree. The order matters: the first column is the outer directory. Put the column you filter on most often first.

PYTHON
# Multi-column: branch (outer) then txn_date (inner)
txn.write.mode("overwrite") \
    .partitionBy("branch", "txn_date") \
    .parquet("/bank/txn-by-branch-date/")
TEXT
Output directory structure:
/bank/txn-by-branch-date/
  branch=Mumbai/
    txn_date=2026-01-15/
      part-00000.parquet
    txn_date=2026-01-16/
      part-00000.parquet
  branch=Delhi/
    txn_date=2026-01-15/
      part-00000.parquet
  branch=Chennai/
    txn_date=2026-01-16/
      part-00000.parquet
PYTHON
# Filter on both: prunes at both levels
result = spark.read.parquet("/bank/txn-by-branch-date/") \
    .filter((F.col("branch") == "Mumbai") & (F.col("txn_date") == "2026-01-16"))
result.explain()
result.show()
TEXT
== Physical Plan ==
FileScan parquet
  PartitionFilters: [branch = Mumbai, txn_date = 2026-01-16]
  ReadSchema: struct<account_id:string,amount:double>

+----------+--------+------+----------+
|account_id|  amount|branch|  txn_date|
+----------+--------+------+----------+
|   SAV2001| 75000.0|Mumbai|2026-01-16|
+----------+--------+------+----------+

What is the Hive-style directory convention?

The column=value directory naming convention comes from Apache Hive. Spark follows this convention by default. The partition column is not stored inside the Parquet files because its value is encoded in the path, which saves a little per-row storage.

PYTHON
# The schema of files inside branch=Mumbai/ does NOT include 'branch'
# Spark reconstructs it from the path when reading
df = spark.read.parquet("/bank/txn-by-branch/")
df.printSchema()
TEXT
root
 |-- account_id: string
 |-- amount: double
 |-- txn_date: string
 |-- branch: string   (reconstructed from directory path)

What goes wrong with high-cardinality partition columns?

Partitioning on a high-cardinality column like account_id or exact timestamp creates millions of tiny directories. This causes directory explosion, the small-file problem, metadata bloat, and can crash the driver during planning.

PYTHON
# BAD: partitioning by account_id (high cardinality)
# If 10 million accounts exist, this creates 10 million directories
# txn.write.partitionBy("account_id").parquet("/bank/bad-idea/")
# Result: 10,000,000 directories, each with a tiny file
# Driver may OOM just listing them

# GOOD: partition by low-cardinality column
txn.write.partitionBy("branch").parquet("/bank/good-idea/")
# Result: ~50 directories (one per branch), each with healthy data

How do you choose the right partition column?

Choose a partition column that meets two criteria: it aligns with how queries actually filter the data, and it has low cardinality so each directory holds a healthy amount of data.

PYTHON
# Decision matrix for choosing partition columns:
# Column         Cardinality   Queried?   Partition?
# branch         ~50           Yes        YES
# country        ~200          Yes        YES
# txn_month      ~12/year      Yes        YES
# account_id     ~10M          Yes        NO (use bucketing)
# exact_ts       ~billions     Rarely     NO
# status         ~5            Sometimes  MAYBE (if selective)

# Good: partition by branch and month
txn.write.partitionBy("branch", "txn_month").parquet("/bank/partitioned/")

# For high-cardinality keys, use bucketing instead
txn.write.bucketBy(64, "account_id").sortBy("txn_date") \
    .saveAsTable("bank.transactions_bucketed")

What is the ideal partition size?

The general guideline is that each partition directory should hold roughly 128 MB to 1 GB of data. This is the range where columnar formats like Parquet read efficiently.

PYTHON
# Check partition sizes after writing
import os

def check_partition_sizes(path):
    for entry in os.listdir(path):
        full = os.path.join(path, entry)
        if os.path.isdir(full) and "=" in entry:
            size = sum(
                os.path.getsize(os.path.join(full, f))
                for f in os.listdir(full)
                if f.endswith(".parquet")
            )
            print(f"  {entry}: {size / 1024 / 1024:.1f} MB")

check_partition_sizes("/bank/txn-by-branch/")
TEXT
  branch=Mumbai: 245.3 MB
  branch=Delhi: 189.7 MB
  branch=Chennai: 312.1 MB

What should a senior engineer know?

A senior engineer keeps the number of partitions in the thousands at most, not millions. They combine partitioning with bucketing when needed: partition by a low-cardinality column for pruning and bucket by a high-cardinality key for shuffle-free joins within each partition.

PYTHON
# Senior pattern: partition + bucket
# Partition by branch (low cardinality, pruning)
# Bucket by account_id (high cardinality, join optimization)
txn.write \
    .partitionBy("branch") \
    .bucketBy(64, "account_id") \
    .sortBy("txn_date") \
    .saveAsTable("bank.transactions_optimized")

# Reading with pruning on branch and shuffle-free join on account_id
optimized = spark.table("bank.transactions_optimized")
mumbai = optimized.filter(F.col("branch") == "Mumbai")  # partition pruning
# Join on account_id is shuffle-free if accounts table is also bucketed

Visual Overview

DIAGRAM
Write-Time Partitioning and Partition Pruning:

WRITE: df.write.partitionBy("branch").parquet(path)
  ┌────────────────────┐
  │   DataFrame        │
  │ account | branch   │
  │ SAV2001 | Mumbai   │───→  /branch=Mumbai/part-0.parquet
  │ SAV2002 | Delhi    │───→  /branch=Delhi/part-0.parquet
  │ CUR3001 | Chennai  │───→  /branch=Chennai/part-0.parquet
  └────────────────────┘

READ with pruning: .filter(branch == "Mumbai")
  /branch=Mumbai/   ← READ (only this directory)
  /branch=Delhi/    ← SKIP
  /branch=Chennai/  ← SKIP

Multi-column partitioning:
  /branch=Mumbai/
    /txn_date=2026-01-15/  part-0.parquet
    /txn_date=2026-01-16/  part-0.parquet
  /branch=Delhi/
    /txn_date=2026-01-15/  part-0.parquet

Cardinality rule:
  LOW cardinality (branch ~50)     → partitionBy  ✓
  HIGH cardinality (account ~10M)  → bucketBy     ✓

Key Takeaway

Write-time partitioning organizes the output physically into a directory hierarchy keyed by a column's values. When I call write with partitionBy on the branch column, Spark creates one subdirectory per branch, named like branch=Mumbai, and it moves the column's value into the directory name rather than storing it in every file. This follows the Hive-style convention of column-equals-value directories.

The payoff comes at read time through partition pruning. When a later query filters on the partition column, the Catalyst Optimizer recognizes that the filter matches the directory structure, so instead of listing and reading every directory it lists only the matching ones and skips the rest entirely. The important point is that this is data skipping at the directory level: the irrelevant data is never read at all, not read and then filtered. That is why pruning is so powerful: a selective filter on a partition column can turn a multi-terabyte scan into a few gigabytes.

In my work at the National Bank of India the transaction lake was partitioned by branch and date, so analysts querying a single branch for a single day touched only that one small directory instead of years of data across all branches. The one caution I always raise is cardinality. Partitioning shines on low-cardinality columns like branch or country where each directory holds a healthy amount of data. If you partition on a high-cardinality column like account_id, you create millions of tiny directories, which bloats metadata, triggers the small-file problem, and can even crash the driver while listing them.

So my rule is: partition by low-cardinality columns you filter on frequently, order the most-filtered column first, and reach for bucketing instead when the key is high-cardinality.

Remember one sentence: partitionBy writes data into value-named directories so queries filter on the partition column skip irrelevant directories entirely, and it works best on low-cardinality columns.

---

Share
Use ← → arrow keys to navigate