fix(rust, python)!: Formalize list aggregation difference between groupbys, selection and window functions

[ritchie46]

This is quite a serious breaking change, but one that has to be made. This was a serious inconsistency in the query engine regarding expressions.

Context

Expressions should adhere to the following rules regarding their context.

1. In the selection context expressions work on columns.
2. In the aggregation/groupby context expressions work on the groups.

Within this you should be able to reason what an expression does. E.g. an expression used in select should explain to you how it should behave in agg as well.

This was not the case.

A col().list() turned a i64 column into a list<i64>.

Based on this, I would assume that an aggregation with col().list() turns a i64 in a list<list<i64>>. The outer list being the groups, the inner list being the list aggregation passed by the user. This was not the case.

But now it is:

Selection

df = pl.DataFrame({
    "group": [1, 2, 2, 3], 
    "value": ["a", "b", "c", "d"]
})

df.select(
    pl.col("value").list()
)

shape: (1, 1)
┌─────────────────────┐
│ value               │
│ ---                 │
│ list[str]           │
╞═════════════════════╡
│ ["a", "b", ... "d"] │
└─────────────────────┘

Groupby

(
    df.groupby("group")
    .agg([
        pl.col("value").alias("values in groups"),
        pl.col("value").list().alias("values in groups + list")
    ])
)

shape: (3, 3)
┌───────┬──────────────────┬─────────────────────────┐
│ group ┆ values in groups ┆ values in groups + list │
│ ---   ┆ ---              ┆ ---                     │
│ i64   ┆ list[str]        ┆ list[list[str]]         │
╞═══════╪══════════════════╪═════════════════════════╡
│ 2     ┆ ["b", "c"]       ┆ [["b", "c"]]            │
│ 1     ┆ ["a"]            ┆ [["a"]]                 │
│ 3     ┆ ["d"]            ┆ [["d"]]                 │
└───────┴──────────────────┴─────────────────────────┘

Window functions

For window functions we must have slightly different mental model. If we don't reduce an aggregation in a window function, e.g. bring down to one element via mean, sum, first etc, we map the groups back to their position in the DataFrame.

For instance a col().sort().over() on a i64 returns the same i64 dtype.
So window function, though very similar to groupby operations, are always a single level less nested than a groupby result.

A col().list().over() will return a list<i64> where the list is the list aggregation and the groups are processed by putting that aggregation on the approriate location in the DataFrame.

This wraps it up. This likely will break some examples in the polars-book so we will have to look at that as well.

Quite breaking, but for the better. :)

Migration

This is easy. Remove any .list() call in a groupby().agg([...]) context.

[stinodego]

In the example you give, shouldn't the groupby result in the values of group 2 each part of an individual list?

shape: (3, 3)
┌───────┬──────────────────┬─────────────────────────┐
│ group ┆ values in groups ┆ values in groups + list │
│ ---   ┆ ---              ┆ ---                     │
│ i64   ┆ list[str]        ┆ list[list[str]]         │
╞═══════╪══════════════════╪═════════════════════════╡
│ 2     ┆ ["b", "c"]       ┆ [["b"], ["c"]]          │
│ 1     ┆ ["a"]            ┆ [["a"]]                 │
│ 3     ┆ ["d"]            ┆ [["d"]]                 │
└───────┴──────────────────┴─────────────────────────┘

Also, interesting change and I agree that it is needed.

[ritchie46]

In the example you give, shouldn't the groupby result in the values of group 2 each part of an individual list?

shape: (3, 3)
┌───────┬──────────────────┬─────────────────────────┐
│ group ┆ values in groups ┆ values in groups + list │
│ ---   ┆ ---              ┆ ---                     │
│ i64   ┆ list[str]        ┆ list[list[str]]         │
╞═══════╪══════════════════╪═════════════════════════╡
│ 2     ┆ ["b", "c"]       ┆ [["b"], ["c"]]          │
│ 1     ┆ ["a"]            ┆ [["a"]]                 │
│ 3     ┆ ["d"]            ┆ [["d"]]                 │
└───────┴──────────────────┴─────────────────────────┘

Also, interesting change and I agree that it is needed.

Nope, that's not what list does in select. You wrap the entire group in a list.

[stinodego]

Nope, that's not what list does in select. You wrap the entire group in a list.

Right, makes sense now that I think about it. Thanks!