Eric Campbell 6 min read Python

PySpark - Partition By in Window Functions

Window functions solve a fundamental limitation in distributed data processing: how do you perform group-based calculations while preserving individual row details? Traditional GROUP BY operations...

Key Insights

  • Window functions with PARTITION BY maintain row-level detail while computing group-based calculations, unlike GROUP BY which collapses rows into aggregates—making them essential for rankings, running totals, and comparative analytics within groups.
  • Partition keys directly impact cluster performance: poor choices create data skew where some executors process significantly more data than others, while well-distributed keys enable parallel processing across the cluster.
  • Combining PARTITION BY with ORDER BY and frame specifications (ROWS BETWEEN/RANGE BETWEEN) unlocks advanced analytics like moving averages and cumulative metrics, but requires understanding how PySpark sorts and windows data within each partition.

Introduction to Window Functions in PySpark

Window functions solve a fundamental limitation in distributed data processing: how do you perform group-based calculations while preserving individual row details? Traditional GROUP BY operations collapse rows into aggregates, losing the granular information needed for rankings, running totals, or row-to-row comparisons.

PySpark’s window functions operate on a “window” of rows related to the current row, allowing you to compute aggregates, rankings, and analytical functions without reducing your dataset. The PARTITION BY clause defines these logical groupings, enabling sophisticated analytics that would otherwise require multiple joins or complex subqueries.

Here’s the fundamental difference:

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import sum, col

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

sales_data = spark.createDataFrame([
    ("East", "Q1", 10000),
    ("East", "Q2", 15000),
    ("West", "Q1", 12000),
    ("West", "Q2", 18000)
], ["region", "quarter", "revenue"])

# GROUP BY: loses row detail
grouped = sales_data.groupBy("region").agg(sum("revenue").alias("total_revenue"))
# Result: 2 rows (East: 25000, West: 30000)

# Window function: maintains all rows
window_spec = Window.partitionBy("region")
windowed = sales_data.withColumn("total_revenue", sum("revenue").over(window_spec))
# Result: 4 rows, each with its quarter revenue AND regional total
windowed.show()

The windowed approach gives you both the individual quarter revenue and the regional total on every row—critical for calculating percentages, rankings, or comparative metrics.

Understanding PARTITION BY Clause

The PARTITION BY clause divides your DataFrame into logical segments based on column values. Each partition is processed independently, which is crucial for distributed computing. When you specify Window.partitionBy("department"), PySpark groups all rows with the same department value together, potentially shuffling data across the cluster to colocate related records.

This partitioning happens at the logical level—it’s not the same as DataFrame partitioning for storage, though the concepts are related. PySpark’s optimizer determines how to physically distribute these logical partitions across executors for parallel processing.

from pyspark.sql.functions import rank, desc

# Sample employee data
employees = spark.createDataFrame([
    ("Engineering", "Alice", 95000),
    ("Engineering", "Bob", 87000),
    ("Engineering", "Charlie", 95000),
    ("Sales", "David", 75000),
    ("Sales", "Eve", 82000),
    ("Sales", "Frank", 78000)
], ["department", "name", "salary"])

# Create window partitioned by department, ordered by salary
dept_window = Window.partitionBy("department").orderBy(desc("salary"))

# Rank employees within their department
ranked_employees = employees.withColumn("rank", rank().over(dept_window))
ranked_employees.show()

# Output shows rankings reset for each department:
# Engineering: Alice=1, Charlie=1, Bob=3
# Sales: Eve=1, Frank=2, David=3

Notice how the rank resets for each department. That’s the power of PARTITION BY—it creates independent contexts for your calculations.

Common Window Function Use Cases with PARTITION BY

Ranking Within Groups

Different ranking functions handle ties differently. ROW_NUMBER always assigns unique sequential numbers, RANK skips numbers after ties, and DENSE_RANK doesn’t skip:

from pyspark.sql.functions import row_number, rank, dense_rank

ranking_window = Window.partitionBy("department").orderBy(desc("salary"))

employees_ranked = employees \
    .withColumn("row_num", row_number().over(ranking_window)) \
    .withColumn("rank", rank().over(ranking_window)) \
    .withColumn("dense_rank", dense_rank().over(ranking_window))

# For tied salaries in Engineering (Alice and Charlie at 95000):
# row_num: 1, 2, 3 (arbitrary order for ties)
# rank: 1, 1, 3 (skips 2)
# dense_rank: 1, 1, 2 (no skip)

Running Totals and Cumulative Aggregations

Running totals require both PARTITION BY and ORDER BY to define the accumulation sequence:

from pyspark.sql.functions import sum

sales = spark.createDataFrame([
    ("Electronics", "2024-01", 50000),
    ("Electronics", "2024-02", 62000),
    ("Electronics", "2024-03", 58000),
    ("Furniture", "2024-01", 35000),
    ("Furniture", "2024-02", 41000),
    ("Furniture", "2024-03", 39000)
], ["category", "month", "sales"])

running_total_window = Window.partitionBy("category").orderBy("month") \
    .rowsBetween(Window.unboundedPreceding, Window.currentRow)

sales_with_running_total = sales.withColumn(
    "running_total", 
    sum("sales").over(running_total_window)
)

# Electronics running total: 50000, 112000, 170000
# Furniture running total: 35000, 76000, 115000

Moving Averages Within Partitions

Moving averages smooth out fluctuations by averaging over a sliding window:

from pyspark.sql.functions import avg

# 3-month moving average per category
moving_avg_window = Window.partitionBy("category").orderBy("month") \
    .rowsBetween(-2, 0)  # Current row and 2 preceding

sales_with_ma = sales.withColumn(
    "moving_avg_3m",
    avg("sales").over(moving_avg_window)
)

# For Electronics month 3: avg(50000, 62000, 58000) = 56666.67

Lead/Lag Operations for Time-Series Analysis

Lead and lag functions access values from subsequent or previous rows within a partition:

from pyspark.sql.functions import lag, lead

comparison_window = Window.partitionBy("category").orderBy("month")

sales_comparison = sales \
    .withColumn("prev_month_sales", lag("sales", 1).over(comparison_window)) \
    .withColumn("next_month_sales", lead("sales", 1).over(comparison_window)) \
    .withColumn("mom_change", col("sales") - lag("sales", 1).over(comparison_window))

# Calculate month-over-month growth within each category

Multiple Column Partitioning

Complex analytics often require partitioning by multiple dimensions. This creates finer-grained groups where calculations are isolated:

# Year-over-year comparison by department and location
detailed_sales = spark.createDataFrame([
    ("Engineering", "NYC", 2023, 500000),
    ("Engineering", "NYC", 2024, 625000),
    ("Engineering", "SF", 2023, 750000),
    ("Engineering", "SF", 2024, 820000),
    ("Sales", "NYC", 2023, 400000),
    ("Sales", "NYC", 2024, 480000)
], ["department", "location", "year", "revenue"])

# Partition by both department AND location
dept_loc_window = Window.partitionBy("department", "location").orderBy("year")

yoy_analysis = detailed_sales.withColumn(
    "prev_year_revenue",
    lag("revenue", 1).over(dept_loc_window)
).withColumn(
    "yoy_growth_pct",
    ((col("revenue") - lag("revenue", 1).over(dept_loc_window)) / 
     lag("revenue", 1).over(dept_loc_window) * 100)
)

# Each department-location combination gets independent year-over-year calculations

The performance trade-off: more partition columns create more granular groups, which can improve parallelism if groups are evenly distributed, but may increase shuffle overhead and create many small partitions if cardinality is high.

Combining PARTITION BY with ORDER BY and Window Frames

The complete window specification includes three components: PARTITION BY (grouping), ORDER BY (ordering within groups), and frame specification (which rows to include in calculations).

Frame specifications use rowsBetween() or rangeBetween():

# rowsBetween: physical row offsets
# rangeBetween: logical value ranges (useful for time-based windows)

time_series = spark.createDataFrame([
    ("Product_A", 1, 100),
    ("Product_A", 2, 150),
    ("Product_A", 3, 120),
    ("Product_A", 5, 180),  # Gap at day 4
    ("Product_A", 6, 200)
], ["product", "day", "quantity"])

# Sum of current row and 2 preceding rows (physical)
rows_window = Window.partitionBy("product").orderBy("day") \
    .rowsBetween(-2, 0)

# Sum within 2-day range (logical - handles gaps)
range_window = Window.partitionBy("product").orderBy("day") \
    .rangeBetween(-2, 0)

result = time_series \
    .withColumn("sum_rows", sum("quantity").over(rows_window)) \
    .withColumn("sum_range", sum("quantity").over(range_window))

# At day 5:
# sum_rows: 120 + 180 + 200 = 500 (3 physical rows)
# sum_range: 180 + 200 = 380 (only days 5 and 6, within 2-day range)

Common frame specifications:

  • Window.unboundedPreceding to Window.currentRow: Running totals
  • Window.unboundedPreceding to Window.unboundedFollowing: Partition-wide aggregates
  • -n to 0: Last n rows including current
  • -n to n: Centered window of 2n+1 rows

Performance Considerations and Best Practices

Window functions trigger shuffle operations to colocate partition data. Poor partition key choices create skewed workloads where some executors process vastly more data than others.

Check partition distribution:

# Examine how data distributes across partition keys
employees.groupBy("department").count().show()

# If one department has 90% of employees, that partition becomes a bottleneck

Repartition strategically:

# If your window partition key differs from DataFrame partitioning
# Repartition first to align physical and logical partitions
optimized_df = employees.repartition("department")

dept_window = Window.partitionBy("department").orderBy("salary")
result = optimized_df.withColumn("rank", rank().over(dept_window))

# This reduces shuffle since data is already colocated by department

Memory management:

Large partitions can cause out-of-memory errors. If a single partition key has millions of rows and you’re using unboundedPreceding to unboundedFollowing, all those rows must fit in executor memory simultaneously.

# Monitor partition sizes
spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

# For very large partitions, consider breaking them up
# or using approximations instead of exact window functions

Avoid unnecessary window recalculation:

# Bad: Creates window spec multiple times
df.withColumn("rank", rank().over(Window.partitionBy("dept").orderBy("salary"))) \
  .withColumn("dense_rank", dense_rank().over(Window.partitionBy("dept").orderBy("salary")))

# Good: Reuse window specification
window = Window.partitionBy("dept").orderBy("salary")
df.withColumn("rank", rank().over(window)) \
  .withColumn("dense_rank", dense_rank().over(window))

Conclusion and Key Takeaways

PARTITION BY in PySpark window functions enables sophisticated analytics that preserve row-level detail while computing group-based metrics. Master these patterns:

Use single-column partitioning for straightforward groupings like rankings by department or running totals by customer. Apply multiple-column partitioning when you need finer granularity, such as year-over-year comparisons by region and product category.

Always consider data distribution when choosing partition keys. High-cardinality keys with even distribution enable better parallelism. Low-cardinality or skewed keys create bottlenecks. Monitor your partition sizes and repartition proactively when your window partition key differs from your DataFrame’s physical partitioning.

Combine PARTITION BY with ORDER BY and frame specifications to unlock the full power of window functions: moving averages with rowsBetween, cumulative sums with unboundedPreceding, and lead/lag operations for time-series analysis.

The most common pitfall is treating window functions like GROUP BY—remember that window functions preserve all rows and add calculated columns rather than reducing your dataset. This makes them perfect for rankings, percentile calculations, and any scenario where you need both individual values and group-level context on the same row.

Liked this? There's more.

Every week: one practical technique, explained simply, with code you can use immediately.