Window Functions#
In this section you will learn about window functions. A window function utilizes values from one or multiple rows to produce a result for each individual row, unlike an aggregate function that provides a single value for multiple rows.
The window functions are available in the functions module.
We’ll use the pokemon dataset (from Ritchie Vink) in the following examples.
from datafusion import SessionContext
from datafusion import col, lit
from datafusion import functions as f
ctx = SessionContext()
df = ctx.read_csv("pokemon.csv")
Here is an example that shows how you can compare each pokemon’s speed to the speed of the previous row in the DataFrame.
df.select(
col('"Name"'),
col('"Speed"'),
f.lag(col('"Speed"')).alias("Previous Speed")
)
DataFrame()
+---------------------------+-------+----------------+
| Name | Speed | Previous Speed |
+---------------------------+-------+----------------+
| Bulbasaur | 45 | |
| Ivysaur | 60 | 45 |
| Venusaur | 80 | 60 |
| VenusaurMega Venusaur | 80 | 80 |
| Charmander | 65 | 80 |
| Charmeleon | 80 | 65 |
| Charizard | 100 | 80 |
| CharizardMega Charizard X | 100 | 100 |
| CharizardMega Charizard Y | 100 | 100 |
| Squirtle | 43 | 100 |
+---------------------------+-------+----------------+
Data truncated.
Setting Parameters#
Ordering#
You can control the order in which rows are processed by window functions by providing
a list of order_by functions for the order_by parameter.
df.select(
col('"Name"'),
col('"Attack"'),
col('"Type 1"'),
f.rank(
partition_by=[col('"Type 1"')],
order_by=[col('"Attack"').sort(ascending=True)],
).alias("rank"),
).sort(col('"Type 1"'), col('"Attack"'))
DataFrame()
+------------+--------+--------+------+
| Name | Attack | Type 1 | rank |
+------------+--------+--------+------+
| Metapod | 20 | Bug | 1 |
| Kakuna | 25 | Bug | 2 |
| Caterpie | 30 | Bug | 3 |
| Weedle | 35 | Bug | 4 |
| Butterfree | 45 | Bug | 5 |
| Venonat | 55 | Bug | 6 |
| Venomoth | 65 | Bug | 7 |
| Paras | 70 | Bug | 8 |
| Beedrill | 90 | Bug | 9 |
| Parasect | 95 | Bug | 10 |
+------------+--------+--------+------+
Data truncated.
Partitions#
A window function can take a list of partition_by columns similar to an
Aggregation Function. This will cause the window values to be evaluated
independently for each of the partitions. In the example above, we found the rank of each
Pokemon per Type 1 partitions. We can see the first couple of each partition if we do
the following:
df.select(
col('"Name"'),
col('"Attack"'),
col('"Type 1"'),
f.rank(
partition_by=[col('"Type 1"')],
order_by=[col('"Attack"').sort(ascending=True)],
).alias("rank"),
).filter(col("rank") < lit(3)).sort(col('"Type 1"'), col("rank"))
DataFrame()
+-----------+--------+----------+------+
| Name | Attack | Type 1 | rank |
+-----------+--------+----------+------+
| Metapod | 20 | Bug | 1 |
| Kakuna | 25 | Bug | 2 |
| Dratini | 64 | Dragon | 1 |
| Dragonair | 84 | Dragon | 2 |
| Voltorb | 30 | Electric | 1 |
| Magnemite | 35 | Electric | 2 |
| Clefairy | 45 | Fairy | 1 |
| Clefable | 70 | Fairy | 2 |
| Mankey | 80 | Fighting | 1 |
| Machop | 80 | Fighting | 1 |
+-----------+--------+----------+------+
Data truncated.
Window Frame#
When using aggregate functions, the Window Frame of defines the rows over which it operates. If you do not specify a Window Frame, the frame will be set depending on the following criteria.
If an
order_byclause is set, the default window frame is defined as the rows between unbounded preceding and the current row.If an
order_byis not set, the default frame is defined as the rows between unbounded and unbounded following (the entire partition).
Window Frames are defined by three parameters: unit type, starting bound, and ending bound.
The unit types available are:
Rows: The starting and ending boundaries are defined by the number of rows relative to the current row.
Range: When using Range, the
order_byclause must have exactly one term. The boundaries are defined bow how close the rows are to the value of the expression in theorder_byparameter.Groups: A “group” is the set of all rows that have equivalent values for all terms in the
order_byclause.
In this example we perform a “rolling average” of the speed of the current Pokemon and the two preceding rows.
from datafusion.expr import Window, WindowFrame
df.select(
col('"Name"'),
col('"Speed"'),
f.avg(col('"Speed"'))
.over(Window(window_frame=WindowFrame("rows", 2, 0), order_by=[col('"Speed"')]))
.alias("Previous Speed"),
)
DataFrame()
+------------+-------+--------------------+
| Name | Speed | Previous Speed |
+------------+-------+--------------------+
| Slowpoke | 15 | 15.0 |
| Jigglypuff | 20 | 17.5 |
| Geodude | 20 | 18.333333333333332 |
| Paras | 25 | 21.666666666666668 |
| Grimer | 25 | 23.333333333333332 |
| Rhyhorn | 25 | 25.0 |
| Snorlax | 30 | 26.666666666666668 |
| Metapod | 30 | 28.333333333333332 |
| Oddish | 30 | 30.0 |
| Parasect | 30 | 30.0 |
+------------+-------+--------------------+
Data truncated.
Chaining onto a window function#
Set a window function’s partition_by, order_by, and window_frame in one
place: its keyword arguments, a single over(), or one builder chain ending in
build(). Chaining a builder method or over() onto a window function that
already has any of them raises:
# Raises: lead already has window options (order_by)
f.lead(col("v"), order_by="t").over(Window(partition_by=[col("g")]))
# Set them together instead.
f.lead(col("v")).over(Window(partition_by=[col("g")], order_by="t"))
A built window function stores a concrete frame with no record of whether you
chose it or it was derived from order_by, so the options cannot be merged
without guessing. Merging may become possible once
apache/datafusion#25934
is resolved.
The null_treatment already set is kept, as are filter and distinct on an
aggregate used as a window function. The whole-partition frame counts as no
frame, so adding an order_by derives the running frame, even when you passed
that frame explicitly:
whole = WindowFrame("rows", None, None) # same as the no-order_by default
# Both give a running sum.
f.sum(col("v")).over(Window()).order_by(col("v")).build()
f.sum(col("v")).over(Window(window_frame=whole)).order_by(col("v")).build()
To keep that frame, set it after the order_by, or pass both in one Window:
f.sum(col("v")).over(Window()).order_by(col("v")).window_frame(whole).build()
f.sum(col("v")).over(Window(order_by=col("v"), window_frame=whole))
Null Treatment#
When using aggregate functions as window functions, it is often useful to specify how null values
should be treated. Pass null_treatment in the Window, or set it on the aggregate itself, which
over() keeps (see Options set on the aggregate).
One common usage for handling nulls is the case where you want to find the last value up to the current row. In the following example we demonstrate how setting the null treatment to ignore nulls will fill in with the value of the most recent non-null row. To do this, we also will set the window frame so that we only process up to the current row.
In this example, we filter down to one specific type of Pokemon that does have some entries in
it’s Type 2 column that are null.
from datafusion.common import NullTreatment
df.filter(col('"Type 1"') == lit("Bug")).select(
'"Name"',
'"Type 2"',
f.last_value(col('"Type 2"'))
.over(
Window(
window_frame=WindowFrame("rows", None, 0),
order_by=[col('"Speed"')],
null_treatment=NullTreatment.IGNORE_NULLS,
)
)
.alias("last_wo_null"),
f.last_value(col('"Type 2"'))
.over(
Window(
window_frame=WindowFrame("rows", None, 0),
order_by=[col('"Speed"')],
null_treatment=NullTreatment.RESPECT_NULLS,
)
)
.alias("last_with_null"),
)
DataFrame()
+------------+--------+--------------+----------------+
| Name | Type 2 | last_wo_null | last_with_null |
+------------+--------+--------------+----------------+
| Paras | Grass | Grass | Grass |
| Metapod | | Grass | |
| Parasect | Grass | Grass | Grass |
| Kakuna | Poison | Poison | Poison |
| Caterpie | | Poison | |
| Venonat | Poison | Poison | Poison |
| Weedle | Poison | Poison | Poison |
| Butterfree | Flying | Flying | Flying |
| Beedrill | Poison | Poison | Poison |
| Pinsir | | Poison | |
+------------+--------+--------------+----------------+
Data truncated.
Aggregate Functions#
You can use any Aggregation Function as a window function. Here
is an example that shows how to compare each pokemons’s attack power with the average attack
power in its "Type 1" using the datafusion.functions.avg() function.
df.select(
col('"Name"'),
col('"Attack"'),
col('"Type 1"'),
f.avg(col('"Attack"')).over(
Window(
window_frame=WindowFrame("rows", None, None),
partition_by=[col('"Type 1"')],
)
).alias("Average Attack"),
)
DataFrame()
+-------------------+--------+--------+----------------+
| Name | Attack | Type 1 | Average Attack |
+-------------------+--------+--------+----------------+
| Dragonite | 134 | Dragon | 94.0 |
| Dratini | 64 | Dragon | 94.0 |
| Dragonair | 84 | Dragon | 94.0 |
| Gastly | 35 | Ghost | 53.75 |
| Haunter | 50 | Ghost | 53.75 |
| Gengar | 65 | Ghost | 53.75 |
| GengarMega Gengar | 65 | Ghost | 53.75 |
| Jynx | 50 | Ice | 67.5 |
| Articuno | 85 | Ice | 67.5 |
| Kabuto | 80 | Rock | 87.5 |
+-------------------+--------+--------+----------------+
Data truncated.
Options set on the aggregate#
over() keeps the filter, distinct, and null_treatment options an
aggregate was built with:
# Averages the distinct values 1.0 and 4.0.
f.avg(col("v"), distinct=True).over(Window())
An aggregate’s order_by raises, as ORDER BY inside an aggregate call does
with OVER in SQL. A window does not pass an ordering to the aggregate, so
the order_by in the Window only sets the frame and the order of rows. The
one exception is a WITHIN GROUP function such as
percentile_cont(), which computes ascending as a
window: an ascending sort_expression is accepted, and a descending one raises.
f.percentile_cont(col("v"), 0.25).over(Window()) # ascending, accepted
f.percentile_cont(col("v").sort(ascending=False), 0.25).over(Window()) # raises
Available Functions#
The possible window functions are:
- Aggregate Functions
All Aggregation Functions can be used as window functions.
User-Defined Window Functions#
You can ship custom window functions to the engine by subclassing
WindowEvaluator and registering it
via udwf(). See datafusion.user_defined
for the evaluator interface and worked examples.
Note
Serialization
Python window UDFs travel inline inside pickled or
to_bytes()-serialized expressions —
the evaluator class is captured by value via cloudpickle, so
worker processes do not need to pre-register the UDF. Any names the
evaluator resolves via import are captured by reference and
must be importable on the receiving worker. See
datafusion.ipc for the full IPC model and security caveats.