Window functions in Apache Spark - Step-by-Step Execution

Concept Flow - Window functions

Start with DataFrame

↓

Define Window Spec

↓

Apply Window Function

↓

Compute Result per Partition

↓

Return DataFrame with New Column

↓

End

Window functions process rows within partitions of data, computing results like ranks or sums without collapsing rows.

Execution Sample

Apache Spark

from pyspark.sql import Window
from pyspark.sql.functions import rank

windowSpec = Window.partitionBy('department').orderBy('salary')
df.withColumn('rank', rank().over(windowSpec)).show()

This code ranks employees by salary within each department.

Execution Table

Step	Action	Input Data	Window Spec	Function Applied	Output Column	Result
1	Start with DataFrame	[{'name':'Alice','department':'HR','salary':3000}, {'name':'Bob','department':'HR','salary':4000}, {'name':'Charlie','department':'IT','salary':3500}]	partitionBy('department').orderBy('salary')	rank()	rank	DataFrame ready for ranking
2	Partition data by department	Same as input	HR: [Alice, Bob], IT: [Charlie]	rank()	rank	Partitions created
3	Order each partition by salary ascending	HR: [Alice(3000), Bob(4000)], IT: [Charlie(3500)]	Same	rank()	rank	Ordered partitions
4	Apply rank() over each partition	Same	Same	rank()	rank	HR: Alice=1, Bob=2; IT: Charlie=1
5	Add rank column to DataFrame	Original rows	Same	rank()	rank	[{'name': 'Alice', 'department': 'HR', 'salary': 3000, 'rank': 1}, {'name': 'Bob', 'department': 'HR', 'salary': 4000, 'rank': 2}, {'name': 'Charlie', 'department': 'IT', 'salary': 3500, 'rank': 1}]
6	Show final DataFrame	Same	Same	rank()	rank	Displayed with rank column
7	End	Final DataFrame	Same	None	None	Execution complete

💡 All rows processed with rank assigned per partition; execution ends.

Variable Tracker

Variable	Start	After Step 2	After Step 3	After Step 4	After Step 5	Final
df	Original DataFrame	Partitioned by department	Partitions ordered by salary	Rank computed per partition	Rank column added	DataFrame with rank column
windowSpec	Not defined	Defined as partitionBy department, orderBy salary	Same	Same	Same	Same
rank	Not defined	Defined as rank() function	Same	Applied over windowSpec	Same	Same

Key Moments - 3 Insights

Why does the rank restart for each department?

Because the windowSpec partitions data by 'department', the rank function resets for each partition as shown in execution_table step 4.

Does the window function reduce the number of rows?

No, window functions add columns but keep all original rows, as seen in execution_table step 5 where rank is added without dropping rows.

What happens if we don't specify orderBy in the windowSpec?

Without orderBy, functions like rank cannot assign meaningful order, so results may be incorrect or error occurs; ordering is essential as shown in step 3.

Visual Quiz - 3 Questions

Look at the execution_table at step 4, what is the rank of Bob in the HR department?

A 1

B 2

C 3

D 0

💡 Hint

Check the 'Result' column at step 4 showing ranks per partition.

At which step is the rank column added to the DataFrame?

A Step 5

B Step 3

C Step 4

D Step 6

💡 Hint

Look for the action 'Add rank column to DataFrame' in the execution_table.

If we remove partitionBy from windowSpec, what changes in the output?

A Rank restarts for each department

B DataFrame rows are reduced

C Rank is computed over entire DataFrame without resetting

D Error occurs immediately

💡 Hint

Partitioning controls where rank resets; without it, rank applies globally as per concept.

Concept Snapshot

Window functions compute values across rows related to the current row.
Use Window specification to define partitions and order.
Apply functions like rank(), row_number(), sum() over the window.
They add columns without reducing rows.
Useful for running totals, rankings, and moving averages.

Full Transcript

Window functions in Apache Spark let you perform calculations across sets of rows related to the current row. You start with a DataFrame, define a window specification that partitions data (like by department) and orders it (like by salary). Then you apply a window function such as rank() over this window. The function computes results per partition and adds a new column to the DataFrame without removing any rows. For example, ranking employees by salary within each department assigns ranks starting at 1 for each department. This process keeps all original data but enriches it with new insights. Key points include that partitioning controls where the function resets, ordering is necessary for ranking, and window functions do not reduce the number of rows.

Practice

1. What is the main purpose of window functions in Apache Spark?

easy

A. To perform calculations across rows related to the current row without collapsing the data

B. To group data and reduce the number of rows

C. To filter rows based on a condition

D. To join two datasets together

Solution

1. Step 1: Understand window function behavior

   Window functions calculate values across related rows but keep all original rows intact.

2. Step 2: Compare with grouping

   Grouping reduces rows by aggregation, but window functions add columns without reducing rows.

3. Final Answer:

   To perform calculations across rows related to the current row without collapsing the data -> Option A

4. Quick Check:

   Window functions add calculations without grouping ≠ A [OK]

Hint: Window functions add columns, grouping reduces rows [OK]

Common Mistakes:

• Confusing window functions with groupBy aggregations
• Thinking window functions filter rows
• Assuming window functions join datasets

2. Which of the following is the correct way to define a window specification in PySpark to partition by 'department' and order by 'salary'?

easy

A. Window.groupBy('department').sortBy('salary')

B. Window.sortBy('salary').partitionBy('department')

C. Window.partitionBy('department').orderBy('salary')

D. Window.partition('department').order('salary')

Solution

1. Step 1: Recall PySpark window spec syntax

   The correct methods are partitionBy() and orderBy() chained on Window.

2. Step 2: Check method order and names

   partitionBy comes first, then orderBy; groupBy and sortBy are incorrect here.

3. Final Answer:

   Window.partitionBy('department').orderBy('salary') -> Option C

4. Quick Check:

   Use partitionBy then orderBy methods [OK]

Hint: Use Window.partitionBy().orderBy() exactly [OK]

Common Mistakes:

• Using groupBy instead of partitionBy
• Using sortBy instead of orderBy
• Using incorrect method names like partition or order

3. Given the following PySpark code, what will be the output of the 'rank' column?

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import rank

spark = SparkSession.builder.getOrCreate()
data = [("Alice", "Sales", 5000), ("Bob", "Sales", 4000), ("Charlie", "Sales", 5000)]
df = spark.createDataFrame(data, ["name", "department", "salary"])
windowSpec = Window.partitionBy("department").orderBy("salary")
df = df.withColumn("rank", rank().over(windowSpec))
df.orderBy("salary").show()

medium

A. [ ("Bob", 4000, 1), ("Alice", 5000, 2), ("Charlie", 5000, 2) ]

B. [ ("Bob", 4000, 1), ("Alice", 5000, 3), ("Charlie", 5000, 3) ]

C. [ ("Bob", 4000, 1), ("Alice", 5000, 2), ("Charlie", 5000, 3) ]

D. [ ("Bob", 4000, 1), ("Alice", 5000, 1), ("Charlie", 5000, 1) ]

Solution

1. Step 1: Understand rank() behavior

   Rank assigns the same rank to ties and skips ranks after ties.

2. Step 2: Apply rank over salary ascending

   Bob has salary 4000 ranked 1; Alice and Charlie both have 5000, so both get rank 2.

3. Final Answer:

   [ ("Bob", 4000, 1), ("Alice", 5000, 2), ("Charlie", 5000, 2) ] -> Option A

4. Quick Check:

   rank() ties share rank, next rank skips [OK]

Hint: rank() gives same rank to ties, skips next ranks [OK]

Common Mistakes:

• Confusing rank() with dense_rank()
• Assigning different ranks to tied values
• Not ordering data before applying rank

4. Identify the error in this PySpark code snippet using window functions:

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number

windowSpec = Window.partitionBy('department')
df = df.withColumn('row_num', row_number().over(windowSpec.orderBy('salary')))

medium

A. Missing import for Window functions causes error

B. windowSpec.orderBy('salary') is valid; no error in code

C. row_number() cannot be used with partitionBy

D. windowSpec.orderBy('salary') is invalid; orderBy should be called on Window, not on windowSpec

Solution

1. Step 1: Check windowSpec usage

   windowSpec is a WindowSpec object; calling orderBy on it is valid and returns a new WindowSpec.

2. Step 2: Confirm chaining

   Chaining orderBy on windowSpec is allowed to refine the window specification.

3. Final Answer:

   windowSpec.orderBy('salary') is valid; no error in code -> Option B

4. Quick Check:

   WindowSpec supports chaining orderBy to refine window [OK]

Hint: WindowSpec supports chaining orderBy to refine window [OK]

Common Mistakes:

• Calling orderBy on windowSpec instead of Window
• Assuming row_number can't be used with partitionBy
• Forgetting to import Window or functions

5. You have a sales dataset with columns: 'store', 'date', and 'sales'. You want to add a column showing the 3-day moving average of sales per store ordered by date. Which PySpark window function setup and code will achieve this?

hard

A. Use Window.partitionBy('store').orderBy('date').rowsBetween(0, 2) and avg('sales').over(windowSpec)

B. Use Window.partitionBy('store').orderBy('date').rangeBetween(-2, 0) and avg('sales').over(windowSpec)

C. Use Window.orderBy('date').rowsBetween(-3, 0) and avg('sales').over(windowSpec)

D. Use Window.partitionBy('store').orderBy('date').rowsBetween(-2, 0) and avg('sales').over(windowSpec)

Solution

1. Step 1: Define window for 3-day moving average per store

   Partition by 'store' to separate stores, order by 'date' for time sequence, use rowsBetween(-2, 0) for current and previous 2 rows.

2. Step 2: Apply avg() over this window

   Use avg('sales').over(windowSpec) to compute moving average over the defined window.

3. Final Answer:

   Use Window.partitionBy('store').orderBy('date').rowsBetween(-2, 0) and avg('sales').over(windowSpec) -> Option D

4. Quick Check:

   rowsBetween(-2,0) for 3-day window, partition by store [OK]

Hint: Use rowsBetween(-2,0) for 3-row moving average [OK]

Common Mistakes:

• Using rangeBetween instead of rowsBetween for row count
• Not partitioning by store
• Using wrong rowsBetween range (like 0 to 2)