Clarify documentation for the agg_list argument in Expr.map_batches

[Wainberg]

Description

The documentation for Expr.map_batches pithily describes the function of the agg_list argument as "Aggregate list". What does this argument do? It would be good to update the documentation.

[deanm0000]

It seems to be for map_elements. I don't think it has a practical usage outside of map_elements using it but I haven't messed with it much

[Wainberg]

map_elements doesn't have an agg_list argument, though.

[deanm0000]

Sorry, to clarify, map_elements calls map_batches and in doing so it sets that parameter in different conditions that I don't remember off hand.

[cmdlineluser]

Looks like it controls ApplyOptions::ApplyList

polars/crates/polars-plan/src/dsl/python_udf.rs

Lines 209 to 210 in a8bdc76

	let (collect_groups, name) = if agg_list {
	(ApplyOptions::ApplyList, MAP_LIST_NAME)

Which is defined here:

polars/crates/polars-plan/src/logical_plan/options.rs

Lines 162 to 172 in a8bdc76

	pub enum ApplyOptions {
	/// Collect groups to a list and apply the function over the groups.
	/// This can be important in aggregation context.
	// e.g. [g1, g1, g2] -> [[g1, g1], g2]
	GroupWise,
	// collect groups to a list and then apply
	// e.g. [g1, g1, g2] -> list([g1, g1, g2])
	ApplyList,
	// do not collect before apply
	// e.g. [g1, g1, g2] -> [g1, g1, g2]
	ElementWise,

[Wainberg]

Does it actually do anything? I haven't been able to find an example where it changes the result.

[reswqa]

They have a clear distinction mainly in the agg context. If agg_list is False, the UDF is called per group. In contrast, the UDF is invoked only once on a list of groups.

Let's use an example to illustrate this further:

df = pl.DataFrame(
         {
            "a": [0,1,0,1],
            "b": [1,2,3,4],
        }
    )

def f(x):
        print(x)
        return x

1. Disable agg_list:

df.group_by("a").agg(pl.col("b").map_batches(f, agg_list=False))

# first output
Series: '' [i64]
[
	2
	4
]

# second output
Series: '' [i64]
[
	1
	3
]

2. Enable agg_list:

df.group_by("a").agg(pl.col("b").map_batches(f, agg_list=True))

# output
Series: 'b' [list[i64]]
[
	[2, 4]
	[1, 3]
]

[reswqa]

Maybe I can update this document to make it easier to understand.