Window functions in Polars offer a versatile way to perform calculations over a range of related rows, making it possible to derive new insights from your data. Unlike simple aggregations or row-level operations, window functions allow you to compute complex metrics across a defined set of rows, known as a window, adding a deeper layer of analysis to your DataFrame.
To understand window functions better, let’s start with a simple dataset:
shape: (5, 3) ┌────────────┬──────────┬───────┐ │ date ┆ category ┆ sales │ │ --- ┆ --- ┆ --- │ │ str ┆ str ┆ i64 │ ╞════════════╪══════════╪═══════╡ │ 2023-01-01 ┆ A ┆ 100 │ │ 2023-01-02 ┆ A ┆ 150 │ │ 2023-01-03 ┆ A ┆ 200 │ │ 2023-01-01 ┆ B ┆ 120 │ │ 2023-01-02 ┆ B ┆ 180 │ └────────────┴──────────┴───────┘
Imagine looking through a window that can move or resize across this data. As this window moves, you can perform calculations on the visible data. For example:
- You might want to calculate a running total of sales over time for each category.
- You could compute a moving average of sales for the past 2 days within each category.
- You might rank the sales within each category.
In each case, the calculation depends not just on the current row, but on a set of related rows defined by your window.
In Polars, the over() method is used to define the window over which these calculations are performed. This approach allows for efficient and flexible data analysis, especially useful in scenarios like time-series analysis, ranking, or calculating running totals and moving averages.
Let’s explore how to use window functions in Polars with some examples using our sample dataset.
Example 1: Basic Window Function
First, let’s create our sample dataset and calculate a cumulative sum using a window function, demonstrating the use of over().
import polars as pl
# Create a sample DataFrame
df = pl.DataFrame({
'date': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-01', '2023-01-02'],
'category': ['A', 'A', 'A', 'B', 'B'],
'sales': [100, 150, 200, 120, 180]
})
# Calculate cumulative sum of sales within each category
df_with_cumsum = df.with_columns([
pl.col('sales').cum_sum().over(['category']).alias('cumulative_sales_by_category')
])
print(df_with_cumsum)
Output:
shape: (5, 4) ┌────────────┬──────────┬───────┬──────────────────────────────┐ │ date ┆ category ┆ sales ┆ cumulative_sales_by_category │ │ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ str ┆ i64 ┆ i64 │ ╞════════════╪══════════╪═══════╪══════════════════════════════╡ │ 2023-01-01 ┆ A ┆ 100 ┆ 100 │ │ 2023-01-02 ┆ A ┆ 150 ┆ 250 │ │ 2023-01-03 ┆ A ┆ 200 ┆ 450 │ │ 2023-01-01 ┆ B ┆ 120 ┆ 120 │ │ 2023-01-02 ┆ B ┆ 180 ┆ 300 │ └────────────┴──────────┴───────┴──────────────────────────────┘
In this example, we use cum_sum() with over([‘category’]) to calculate the cumulative sum of sales for each category independently. Notice how the cumulative sum restarts for each category.
Example 2: Moving Average
Now, let’s calculate a 2-day moving average of sales within each category.
# Calculate 2-day moving average within each category
df_with_ma = df.with_columns([
pl.col('sales').rolling_mean(window_size=2).over(['category']).alias('2day_moving_avg')
])
print(df_with_ma)
Output:
shape: (5, 4) ┌────────────┬──────────┬───────┬─────────────────┐ │ date ┆ category ┆ sales ┆ 2day_moving_avg │ │ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ str ┆ i64 ┆ f64 │ ╞════════════╪══════════╪═══════╪═════════════════╡ │ 2023-01-01 ┆ A ┆ 100 ┆ null │ │ 2023-01-02 ┆ A ┆ 150 ┆ 125.0 │ │ 2023-01-03 ┆ A ┆ 200 ┆ 175.0 │ │ 2023-01-01 ┆ B ┆ 120 ┆ null │ │ 2023-01-02 ┆ B ┆ 180 ┆ 150.0 │ └────────────┴──────────┴───────┴─────────────────┘
Here, rolling_mean(window_size=2) is used with over([‘category’]) to compute the moving average over a 2-day window for each category separately. Note that the first day in each category has a null value as there’s not enough data for a 2-day average.
Example 3: Ranking
Let’s rank the sales within each category.
# Rank sales within each category
df_with_rank = df.with_columns([
pl.col('sales').rank().over('category').alias('sales_rank_in_category')
])
print(df_with_rank)
Output:
shape: (5, 4) ┌────────────┬──────────┬───────┬────────────────────────┐ │ date ┆ category ┆ sales ┆ sales_rank_in_category │ │ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ str ┆ i64 ┆ f64 │ ╞════════════╪══════════╪═══════╪════════════════════════╡ │ 2023-01-01 ┆ A ┆ 100 ┆ 1.0 │ │ 2023-01-02 ┆ A ┆ 150 ┆ 2.0 │ │ 2023-01-03 ┆ A ┆ 200 ┆ 3.0 │ │ 2023-01-01 ┆ B ┆ 120 ┆ 1.0 │ │ 2023-01-02 ┆ B ┆ 180 ┆ 2.0 │ └────────────┴──────────┴───────┴────────────────────────┘
In this example, we use rank() with over(‘category’) to rank sales within each category independently. The sales are ranked from lowest (1) to highest within each category.
Example 4: Advanced Window Function
Finally, let’s perform a more complex operation: calculating the percentage of total sales for each day within its category.
# Calculate percentage of total sales for each day within its category
df_advanced = df.with_columns([
(pl.col('sales') / pl.col('sales').sum().over('category') * 100).alias('percent_of_category_sales')
])
print(df_advanced)
Output:
shape: (5, 4) ┌────────────┬──────────┬───────┬───────────────────────────┐ │ date ┆ category ┆ sales ┆ percent_of_category_sales │ │ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ str ┆ i64 ┆ f64 │ ╞════════════╪══════════╪═══════╪═══════════════════════════╡ │ 2023-01-01 ┆ A ┆ 100 ┆ 22.222222 │ │ 2023-01-02 ┆ A ┆ 150 ┆ 33.333333 │ │ 2023-01-03 ┆ A ┆ 200 ┆ 44.444444 │ │ 2023-01-01 ┆ B ┆ 120 ┆ 40.0 │ │ 2023-01-02 ┆ B ┆ 180 ┆ 60.0 │ └────────────┴──────────┴───────┴───────────────────────────┘
Here, we divide each day’s sales by the total sales for its category (calculated using sum().over(‘category’)), then multiply by 100 to get the percentage.
Conclusion
Window functions in Polars, implemented using the over() method, provide an efficient way to perform complex calculations across related rows of data. They’re particularly useful for time-series analysis, rankings, and calculating running totals or averages within specific groups or categories. By practicing these functions, you can efficiently perform advanced data manipulations and gain deeper insights from your datasets.
