How can I calculate running totals using PySpark window functions?

Below is an example of how to use PySpark window functions. Window functions allow you to perform operations (such as ranking, running totals, and moving averages) over a group of rows that are somehow related to the current row.

Example Scenario

Suppose we have a DataFrame with sales data and we want to calculate a running total of sales per store, ordered by date.

Step-by-Step Example

1. Import the required modules:

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

2. Create a Spark session and sample DataFrame:

PYTHONCopied
spark = SparkSession.builder.appName("WindowFunctionExample").getOrCreate()

data = [
    ("store_1", "2023-01-01", 100),
    ("store_1", "2023-01-02", 150),
    ("store_1", "2023-01-03", 200),
    ("store_2", "2023-01-01", 80),
    ("store_2", "2023-01-02", 120),
    ("store_2", "2023-01-03", 160)
]

columns = ["store", "date", "sales"]

df = spark.createDataFrame(data, columns)

3. Define the window specification:

The window defines how the data is partitioned (grouped) and ordered.

PYTHONCopied
window_spec = Window.partitionBy("store").orderBy("date") \
                    .rowsBetween(Window.unboundedPreceding, Window.currentRow)

4. Apply the window function:

In this example, we apply a running total (cumulative sum) on the sales column.

PYTHONCopied
df_cumulative = df.withColumn("running_total", spark_sum(col("sales")).over(window_spec))
df_cumulative.show()

Expected Output

The result will display the running total for each store:

+-------+----------+-----+-------------+
| store | date     |sales|running_total|
+-------+----------+-----+-------------+
|store_1|2023-01-01| 100 |         100 |
|store_1|2023-01-02| 150 |         250 |
|store_1|2023-01-03| 200 |         450 |
|store_2|2023-01-01| 80  |          80 |
|store_2|2023-01-02| 120 |         200 |
|store_2|2023-01-03| 160 |         360 |
+-------+----------+-----+-------------+

Additional Notes

• Partitioning:
  Partitioning by "store" means the window function resets its calculation for every store.

• Ordering:
  Ordering by "date" ensures that the cumulative sum follows the chronological order.

• Frame Specification:
  The frame specification (rowsBetween(Window.unboundedPreceding, Window.currentRow)) calculates the cumulative sum from the first row in the partition up to the current row.

This approach is common for tasks like calculating running totals, rankings, or moving averages in PySpark.