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 functionality of window functions in DataFusion is supported by the dedicated window() function.

We’ll use the pokemon dataset (from Ritchie Vink) in the following examples.

In [1]: import urllib.request

In [2]: from datafusion import SessionContext

In [3]: from datafusion import col

In [4]: from datafusion import functions as f

In [5]: urllib.request.urlretrieve(
   ...:     "https://gist.githubusercontent.com/ritchie46/cac6b337ea52281aa23c049250a4ff03/raw/89a957ff3919d90e6ef2d34235e6bf22304f3366/pokemon.csv",
   ...:     "pokemon.csv",
   ...: )
   ...: 
Out[5]: ('pokemon.csv', <http.client.HTTPMessage at 0x7f9df0dc4550>)

In [6]: ctx = SessionContext()

In [7]: df = ctx.read_csv("pokemon.csv")

Here is an example that shows how to compare each pokemons’s attack power with the average attack power in its "Type 1"

In [8]: df.select(
   ...:     col('"Name"'),
   ...:     col('"Attack"'),
   ...:     f.alias(
   ...:         f.window("avg", [col('"Attack"')], partition_by=[col('"Type 1"')]),
   ...:         "Average Attack",
   ...:     )
   ...: )
   ...: 
Out[8]: 
DataFrame()
+-----------------------+--------+--------------------+
| Name                  | Attack | Average Attack     |
+-----------------------+--------+--------------------+
| Exeggcute             | 40     | 40.0               |
| Ivysaur               | 62     | 51.0               |
| Venusaur              | 82     | 61.333333333333336 |
| VenusaurMega Venusaur | 100    | 71.0               |
| Tangela               | 55     | 67.8               |
| Bellsprout            | 75     | 69.0               |
| Vileplume             | 80     | 70.57142857142857  |
| Gloom                 | 65     | 69.875             |
| Bulbasaur             | 49     | 67.55555555555556  |
| Oddish                | 50     | 65.8               |
+-----------------------+--------+--------------------+

You can also control the order in which rows are processed by window functions by providing a list of order_by() functions for the order_by parameter.

In [9]: df.select(
   ...:     col('"Name"'),
   ...:     col('"Attack"'),
   ...:     f.alias(
   ...:         f.window(
   ...:             "rank",
   ...:             [],
   ...:             partition_by=[col('"Type 1"')],
   ...:             order_by=[f.order_by(col('"Attack"'))],
   ...:         ),
   ...:         "rank",
   ...:     ),
   ...: )
   ...: 
Out[9]: 
DataFrame()
+------------+--------+------+
| Name       | Attack | rank |
+------------+--------+------+
| Dratini    | 64     | 1    |
| Dragonair  | 84     | 2    |
| Dragonite  | 134    | 3    |
| Clefairy   | 45     | 1    |
| Clefable   | 70     | 2    |
| Mankey     | 80     | 1    |
| Machop     | 80     | 1    |
| Machoke    | 100    | 3    |
| Primeape   | 105    | 4    |
| Hitmonchan | 105    | 4    |
+------------+--------+------+

The possible window functions are:

  1. Rank Functions
    • rank

    • dense_rank

    • row_number

    • ntile

  2. Analytical Functions
    • cume_dist

    • percent_rank

    • lag

    • lead

    • first_value

    • last_value

    • nth_value

  3. Aggregate Functions
    • All aggregate functions can be used as window functions.