Unlocking Insights with Spark SQL Window Functions#

Ranking Functions#

Ranking functions are useful for assigning row numbers or ranks to records, often for the purpose of ordering or identifying top records within each partition.

Aggregate Functions#

Aggregate functions can be applied in a window to compute calculations like sums, averages, counts, etc., without collapsing rows.

Analytical Functions#

Analytical functions operate on the values in each row relative to other rows in the partition.

Rows Between vs. Range Between#

• ROWS BETWEEN: The bounds define specific numeric offsets from the current row. For example, ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING includes exactly three rows in the window: the previous row, the current row, and the next row.
• RANGE BETWEEN: The bounds define a logical range of values relative to the current row’s ordering. For example, RANGE BETWEEN INTERVAL 7 DAYS PRECEDING AND CURRENT ROW might include all rows in the past week up to the current row’s timestamp (in a time-ordered dataset).

Unbounded vs. Bounded Frames#

• UNBOUNDED PRECEDING: Points to the start of the partition.
• UNBOUNDED FOLLOWING: Points to the end of the partition.
• CURRENT ROW: Points to the current row’s position.

A common case:

1
SUM(sales) OVER (
2
  PARTITION BY region
3
  ORDER BY sales_date
4
  ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
5
) AS cumulative_sales

This calculates a cumulative sum of sales over time within each region partition.

Example Dataset Setup#

Imagine we have a table named orders with the following schema:

For demonstration, we can create or register a temporary view in Spark SQL. If you already have such a table in a Hive metastore, skip this step. Otherwise:

1
from pyspark.sql import Row
2
from datetime import datetime
3

4
data = [
5
    Row(order_id=1, customer_id=101, order_date=datetime(2023, 1, 1), order_amount=100.0, category="Electronics"),
6
    Row(order_id=2, customer_id=101, order_date=datetime(2023, 1, 3), order_amount=80.0,  category="Electronics"),
7
    Row(order_id=3, customer_id=102, order_date=datetime(2023, 1, 2), order_amount=50.0,  category="Home"),
8
    Row(order_id=4, customer_id=102, order_date=datetime(2023, 1, 5), order_amount=120.0, category="Home"),
9
    Row(order_id=5, customer_id=103, order_date=datetime(2023, 1, 2), order_amount=60.0,  category="Clothing"),
10
    Row(order_id=6, customer_id=103, order_date=datetime(2023, 1, 4), order_amount=200.0, category="Clothing"),
11
    Row(order_id=7, customer_id=104, order_date=datetime(2023, 1, 6), order_amount=150.0, category="Electronics")
12
]
13

14
df = spark.createDataFrame(data)
15
df.createOrReplaceTempView("orders")

Now we can use Spark SQL queries directly:

1
spark.sql("SELECT * FROM orders").show()

Ranking with ROW_NUMBER()#

A simple yet powerful use case is ranking row entries within a partition. Suppose we want to find the top 3 orders by order_amount within each category:

1
SELECT
2
  category,
3
  order_id,
4
  order_amount,
5
  ROW_NUMBER() OVER (
6
    PARTITION BY category
7
    ORDER BY order_amount DESC
8
  ) AS order_rank
9
FROM orders

From the result, each category partition is separated, and within each category, ordering is by descending order_amount. The highest order_amount in each category will be assigned 1, the next highest 2, and so forth.

Rolling Averages#

If we want to determine, for each customer, the rolling average of their order amounts over all time, we can specify a frame that includes all previous orders up to the current order (unbounded preceding to current row):

1
SELECT
2
  customer_id,
3
  order_date,
4
  order_amount,
5
  AVG(order_amount) OVER (
6
    PARTITION BY customer_id
7
    ORDER BY order_date
8
    ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
9
  ) AS rolling_avg
10
FROM orders
11
ORDER BY customer_id, order_date;

This query calculates a cumulative average up to each order date per customer.

Cumulative Sums#

Cumulative sums are another popular scenario. If we want to see how much each customer has spent up to each particular order date, simply replace AVG with SUM:

1
SELECT
2
  customer_id,
3
  order_date,
4
  order_amount,
5
  SUM(order_amount) OVER (
6
    PARTITION BY customer_id
7
    ORDER BY order_date
8
    ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
9
  ) AS running_total
10
FROM orders
11
ORDER BY customer_id, order_date;

Lag and Lead Operations#

For analyses that require comparing a row to a previous or next row, LAG and LEAD are essential. One practical use case is finding the difference in order_amount between the current order and the previous order for each customer:

1
SELECT
2
  customer_id,
3
  order_date,
4
  order_amount,
5
  LAG(order_amount, 1, 0) OVER (
6
    PARTITION BY customer_id
7
    ORDER BY order_date
8
  ) AS prev_amount,
9
  (order_amount - LAG(order_amount, 1, 0) OVER (
10
    PARTITION BY customer_id
11
    ORDER BY order_date
12
  )) AS diff_from_previous
13
FROM orders
14
ORDER BY customer_id, order_date;

Here we are using LAG(order_amount, 1, 0) which returns 0 if there is no previous row for that customer (the default value). Then we compute the difference for the current row.

Combining Multiple Window Functions#

In many scenarios, you might need to break down a dataset using different partitioning or ordering criteria for multiple metrics. One approach is using subqueries or CTEs:

1
WITH base AS (
2
  SELECT
3
    order_id,
4
    customer_id,
5
    category,
6
    order_amount,
7
    RANK() OVER (PARTITION BY category ORDER BY order_amount DESC) AS category_rank,
8
    SUM(order_amount) OVER (
9
      PARTITION BY category
10
      ORDER BY order_date
11
      ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
12
    ) AS category_cumulative
13
  FROM orders
14
)
15
SELECT
16
  order_id,
17
  customer_id,
18
  category,
19
  order_amount,
20
  category_rank,
21
  category_cumulative
22
FROM base
23
WHERE category_rank <= 3
24
ORDER BY category, category_rank;

In this example, we first define a CTE called base that calculates two window functions at once: a rank by category and a cumulative sum by category. The outer query can then filter or further manipulate these calculations.

Performance Tuning#

Window functions, while powerful, can be resource-intensive. Remember the following tips:

1. Partition Selection: Narrow partitions based on relevant columns. This helps Spark distribute computations more efficiently.
2. Ordering: Ensure the ordering field is properly indexed or well-partitioned if possible. In many big-data contexts, physically sorting data can be expensive.
3. Frame Definitions: Use tighter frame definitions. Instead of using UNBOUNDED PRECEDING, choose smaller frames if your analysis allows. This can reduce the amount of data Spark reads per row.
4. Cluster Configuration: Adjust the number of shuffle partitions (spark.sql.shuffle.partitions) if you see skew in partition distribution.
5. Predicate Pushdown: Filter datasets before using window functions whenever possible, to avoid unnecessarily large data movements.

Memory Management#

Window functions require data shuffling across partitions, often leading to high disk and memory usage. Large frames can cause Spark to spill data to disk. Ensure the following:

• Your Spark executors have enough memory overhead.
• Consider using persistent storage when intermediate results are repeatedly accessed.
• Evaluate the data distribution. Data skew can lead to out-of-memory errors in certain partitions.

Complex Frame Definitions#

If your logic requires, you can define more intricate frames. For instance, if you want a 7-day rolling window for sales:

1
SELECT
2
  order_date,
3
  order_amount,
4
  SUM(order_amount) OVER (
5
    ORDER BY order_date
6
    RANGE BETWEEN 7 PRECEDING AND CURRENT ROW
7
  ) AS rolling_7_day_sales
8
FROM orders
9
ORDER BY order_date;

This approach might utilize a timestamp or date ordering column and define an interval that captures 7 days. Spark interprets 7 PRECEDING differently under RANGE—it’s typically valid for numeric columns or intervals if you are using a date/timestamp type with certain syntax.

Nested Window Functions#

Nested window functions go beyond typical usage, such as performing a window function on top of a subquery that already contains a window function. This can be beneficial for complex analytics like computing a moving average of a moving average. However, be mindful of performance implications and complexity.

An example of a two-step process might involve a subquery that calculates a daily rolling average, with an outer query that ranks these daily rolling averages:

1
WITH daily_avg AS (
2
  SELECT
3
    order_date,
4
    customer_id,
5
    AVG(order_amount) OVER (
6
      PARTITION BY customer_id
7
      ORDER BY order_date
8
      ROWS BETWEEN 2 PRECEDING AND CURRENT ROW
9
    ) AS rolling_avg_3d
10
  FROM orders
11
)
12
SELECT
13
  customer_id,
14
  order_date,
15
  rolling_avg_3d,
16
  RANK() OVER (ORDER BY rolling_avg_3d DESC) AS rolling_avg_rank
17
FROM daily_avg;

Use Cases in Data Science and Machine Learning#

Window functions also play a complementary role in data science workflows:

1. Feature Engineering: Create new features based on rolling stats or cumulative sums for machine learning models.
2. Time Series Analysis: Calculate lagged features (e.g., previous day sales, previous week’s average) for forecasting models.
3. Anomaly Detection: Use window functions to detect outliers by comparing each value to rolling averages or medians.

When integrated with a broader Spark ML pipeline, these transformations become an integral part of your data preparation steps.