docs: Explain differences between `list.eval` and `list.agg`

[cr7pt0gr4ph7]

Description

There is a dedicated list.agg method for aggregations, as well as an additional list.eval method that can do both element-wise mapping as well as aggregations. It's currently not clear when to use which - even the linked documentation of the two methods only contains examples for which the two yield the same results when replaced with the other - or what advantage one has over the other.

Even the docstring for list.agg presents the following example, which arguably isn't an aggregation (meaning it maps many values to a single value) but a per-element filter:

>> df.with_columns(no_nulls=pl.col.a.list.agg(pl.element().drop_nulls()))
shape: (3, 2)
┌──────────────┬───────────┐
│ a            ┆ no_nulls  │
│ ---          ┆ ---       │
│ list[i64]    ┆ list[i64] │
╞══════════════╪═══════════╡
│ [1, null]    ┆ [1]       │
│ [42, 13]     ┆ [42, 13]  │
│ [null, null] ┆ []        │
└──────────────┴───────────┘

...so it arguably would probably be more accurate to use list.eval instead:

>>> df.with_columns(no_nulls=pl.col.a.list.agg(pl.element().drop_nulls()))
shape: (3, 2)
┌──────────────┬───────────┐
│ a            ┆ no_nulls  │
│ ---          ┆ ---       │
│ list[i64]    ┆ list[i64] │
╞══════════════╪═══════════╡
│ [1, null]    ┆ [1]       │
│ [42, 13]     ┆ [42, 13]  │
│ [null, null] ┆ []        │
└──────────────┴───────────┘

The only obvious difference I can see is that list.eval has a parallel=True|False parameter, while list.agg doesn't. The confusion is compounded by the fact that for top-level mapping or aggregation, there is only a single DataFrame.select method instead of two separate agg/eval methods.

Link

• https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.list.agg.html
• https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.list.eval.html

[cr7pt0gr4ph7]

This issue spawned from this review comment by @coastalwhite over here:

• docs: Explain aggregation & sorting of lists #25260

[cmdlineluser]

I believe agg was added so you don't need to use .list.eval(...).list.first() or so when reducing a list down to a single item:

pl.select(x=[1, 2, 3]).with_columns(
    eval = pl.col.x.list.eval(pl.element().entropy()),
    agg = pl.col.x.list.agg(pl.element().entropy())
)
# shape: (1, 3)
# ┌───────────┬────────────┬──────────┐
# │ x         ┆ eval       ┆ agg      │
# │ ---       ┆ ---        ┆ ---      │
# │ list[i64] ┆ list[f64]  ┆ f64      │
# ╞═══════════╪════════════╪══════════╡
# │ [1, 2, 3] ┆ [1.011404] ┆ 1.011404 │
# └───────────┴────────────┴──────────┘

I'm not sure if parallel= does anything anymore? I believe it was used in early implementations of list.eval before it was changed to run on the group_by engine - not 100% sure though.

[cr7pt0gr4ph7]

I'm currently digging through the code, I'll get back shortly when I can list their differences (if any).

[cr7pt0gr4ph7]

After digging into the code:

• They both map to an Expr::Eval node, only differing in EvalVariant::List vs. EvalVariant::ListAgg:

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

Lines 300 to 314 in 82f9741

	pub fn eval<E: Into<Expr>>(self, other: E) -> Expr {
	Expr::Eval {
	expr: Arc::new(self.0),
	evaluation: Arc::new(other.into()),
	variant: EvalVariant::List,
	}
	}
	pub fn agg<E: Into<Expr>>(self, other: E) -> Expr {
	Expr::Eval {
	expr: Arc::new(self.0),
	evaluation: Arc::new(other.into()),
	variant: EvalVariant::ListAgg,
	}
	}

---

Searching for EvalVariant::List/EvalVariant::ListAgg yields the following behavioral differences and similarities:

• Formatted as list.agg(...) vs. list.eval(...) when formatting as a string:

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

Lines 194 to 195 in 82f9741

	EvalVariant::List => write!(f, "{input:?}.list.eval({evaluation:?})"),
	EvalVariant::ListAgg => write!(f, "{input:?}.list.agg({evaluation:?})"),

• Treated the same for purposes of expression pushdown by the optimizer:

polars/crates/polars-plan/src/plans/optimizer/set_order/expr_pushdown.rs

Lines 293 to 310 in 82f9741

	AExpr::Eval {
	expr,
	evaluation: _,
	variant,
	} => match variant {
	EvalVariant::Array { as_list: _ }
	| EvalVariant::ArrayAgg
	| EvalVariant::List
	| EvalVariant::ListAgg => rec!(*expr),
	EvalVariant::Cumulative { min_samples: _ } => {
	let expr = rec!(*expr);
	if expr.column_ordering_observable() {
	return Err(ColumnOrderObserved);
	}
	expr
	},
	},

• Have the same dtype for pl.element():

polars/crates/polars-plan/src/dsl/expr/mod.rs

Lines 604 to 612 in 82f9741

	pub fn element_dtype<'a>(&self, dtype: &'a DataType) -> PolarsResult<&'a DataType> {
	match (self, dtype) {
	(Self::List | Self::ListAgg, DataType::List(inner)) => Ok(inner.as_ref()),
	#[cfg(feature = "dtype-array")]
	(Self::Array { .. } | Self::ArrayAgg, DataType::Array(inner, _)) => Ok(inner.as_ref()),
	(Self::Cumulative { min_samples: _ }, dt) => Ok(dt),
	_ => polars_bail!(op = self.to_name(), dtype),
	}
	}

• Both are lowered to AExpr::Eval in the streaming engine, with the EvalVariant being passed through:

polars/crates/polars-stream/src/physical_plan/lower_expr.rs

Lines 1500 to 1512 in 82f9741

	EvalVariant::List
	| EvalVariant::ListAgg
	| EvalVariant::Array { as_list: _ }
	| EvalVariant::ArrayAgg => {
	let (trans_input, trans_expr) = lower_exprs_with_ctx(input, &[inner], ctx)?;
	let eval_expr = AExpr::Eval {
	expr: trans_expr[0],
	evaluation,
	variant,
	};
	input_streams.insert(trans_input);
	transformed_exprs.push(ctx.expr_arena.add(eval_expr));
	},

• In the alternative expr_to_ir conversion, they are also treated exactly the same, with the EvalVariant being passed through:

polars/crates/polars-plan/src/plans/conversion/dsl_to_ir/expr_to_ir.rs

Lines 445 to 500 in 82f9741

	Expr::Eval {
	expr,
	evaluation,
	variant,
	} => {
	let (expr, output_name) = recurse_arc!(expr)?;
	let expr_dtype = ctx.arena.get(expr).to_dtype(&ctx.to_field_ctx())?;
	let element_dtype = variant.element_dtype(&expr_dtype)?;
	// Perform this before schema resolution so that we can better error messages.
	for e in evaluation.as_ref().into_iter() {
	if matches!(e, Expr::Column(_)) {
	polars_bail!(
	ComputeError:
	"named columns are not allowed in `eval` functions; consider using `element`"
	);
	}
	}
	let mut evaluation_schema = ctx.schema.clone();
	evaluation_schema.insert(PL_ELEMENT_NAME.clone(), element_dtype.clone());
	let mut evaluation_ctx = ExprToIRContext {
	with_fields: None,
	schema: &evaluation_schema,
	arena: ctx.arena,
	allow_unknown: ctx.allow_unknown,
	check_column_names: ctx.check_column_names,
	};
	let (evaluation, _) = to_aexpr_impl(owned(evaluation), &mut evaluation_ctx)?;
	match variant {
	EvalVariant::List | EvalVariant::ListAgg => {},
	EvalVariant::Array { as_list } => {
	polars_ensure!(
	as_list || is_length_preserving_ae(evaluation, ctx.arena),
	InvalidOperation: "`array.eval` is not allowed with non-length preserving expressions. Enable `as_list` if you want to output a variable amount of items per row."
	)
	},
	EvalVariant::ArrayAgg => {},
	EvalVariant::Cumulative { .. } => {
	polars_ensure!(
	is_scalar_ae(evaluation, ctx.arena),
	InvalidOperation: "`cumulative_eval` is not allowed with non-scalar output"
	)
	},
	}
	(
	AExpr::Eval {
	expr,
	evaluation,
	variant,
	},
	output_name,
	)
	},

What we can see here is that arr.agg and arr.eval do have some behavioral differences, as arr.eval is normally required to be length preserving (except when as_list=True is specified). list.agg and list.eval were probably created to just mirror the arr API.

• Both are mapped to the physical EvalExpr implementation, with the EvalVariant being passed through:

polars/crates/polars-expr/src/planner.rs

Lines 584 to 594 in 82f9741

	Ok(Arc::new(EvalExpr::new(
	expr,
	evaluation,
	*variant,
	node_to_expr(expression, expr_arena),
	output_field,
	is_scalar,
	evaluation_is_scalar,
	evaluation_is_elementwise,
	evaluation_is_fallible,
	)))

What's interesting here is a code section slightly above that independently determines whether we have an aggregation expression (evaluation_is_scalar) or an element-wise expression (evaluation_is_elementwise):

polars/crates/polars-expr/src/planner.rs

Lines 546 to 552 in 82f9741

	let is_scalar = is_scalar_ae(expression, expr_arena);
	let evaluation_is_scalar = is_scalar_ae(*evaluation, expr_arena);
	let evaluation_is_elementwise = is_elementwise_rec(*evaluation, expr_arena);
	// @NOTE: This is actually also something the downstream apply code should care about.
	let mut pd_group = ExprPushdownGroup::Pushable;
	pd_group.update_with_expr_rec(expr_arena.get(*evaluation), expr_arena, None);
	let evaluation_is_fallible = matches!(pd_group, ExprPushdownGroup::Fallible);

• The physical expression implementation treats them nearly the same, except for an ìs_agg` parameter:

polars/crates/polars-expr/src/expressions/eval.rs

Lines 412 to 419 in 82f9741

	EvalVariant::List => {
	let lst = input.list()?;
	self.evaluate_on_list_chunked(lst, state, false)
	},
	EvalVariant::ListAgg => {
	let lst = input.list()?;
	self.evaluate_on_list_chunked(lst, state, true)
	},

[cr7pt0gr4ph7]

The only differences therefore seem to be in the evaluate_on_list_chunked function. Note that the is_agg parameter is never checked by itself, but only in combination with expression_is_elementwise and expression_is_scalar:

polars/crates/polars-expr/src/expressions/eval.rs

Lines 84 to 100 in 82f9741

	// Fast path: fully elementwise expression without masked out values.
	if self.evaluation_is_elementwise && !may_fail_on_masked_out_elements {
	let mut column = self.evaluation.evaluate(&df, state)?;
	// Since `lit` is marked as elementwise, this may lead to problems.
	if column.len() == 1 && df.height() != 1 {
	column = column.new_from_index(0, df.height());
	}
	if !is_agg || !self.evaluation_is_scalar {
	column = ca
	.with_inner_values(column.as_materialized_series())
	.into_column();
	}
	return Ok(column);
	}

polars/crates/polars-expr/src/expressions/eval.rs

Lines 133 to 156 in 82f9741

	// Fast path. Groups are pointing to the same offsets in the data buffer.
	if flat_naive.len() == df.height()
	&& let Some(output_groups) = ac.groups.as_ref().as_unrolled_slice()
	&& !(is_agg && self.evaluation_is_scalar)
	{
	let groups_are_unchanged = if let Some(validity) = &validity {
	assert_eq!(validity.set_bits(), output_groups.len());
	validity
	.true_idx_iter()
	.zip(output_groups)
	.all(|(j, [start, len])| {
	let (original_start, original_end) =
	unsafe { offsets.start_end_unchecked(j) };
	(*start == original_start as IdxSize)
	& (*len == (original_end - original_start) as IdxSize)
	})
	} else {
	output_groups
	.iter()
	.zip(offsets.offset_and_length_iter())
	.all(|([start, len], (original_start, original_len))| {
	(*start == original_start as IdxSize) & (*len == original_len as IdxSize)
	})
	};

polars/crates/polars-expr/src/expressions/eval.rs

Lines 164 to 183 in 82f9741

	// Slow path. Groups have changed, so we need to gather data again.
	if is_agg && self.evaluation_is_scalar {
	let mut values = ac.finalize();
	// We didn't have any groups for the `null` values so we have to reinsert them.
	if let Some(validity) = validity {
	values = values.deposit(&validity);
	}
	Ok(values)
	} else {
	let mut ca = ac.aggregated_as_list();
	// We didn't have any groups for the `null` values so we have to reinsert them.
	if let Some(validity) = validity {
	ca = Cow::Owned(ca.deposit(&validity));
	}
	Ok(ca.into_owned().into_column())
	}

[coastalwhite]

The list.agg and list.eval expressions are exactly the same, except one difference. If the evaluation expression is statically determined to return only one value, it will automatically explode the list into the inner values. This matches what .group_by(...).agg(...) does, hence the name.

In my opinion, list.eval has a tiny amount of use cases left, but it is mostly there for backwards compatibility.

The situation is somewhat similar for arr.eval and arr.agg, but there is one slight difference there. arr.eval is able to keep the array if the evaluation expression is statically determined to be length preserving.

Those should be the only differences as the implementations are almost entirely shared.

[cr7pt0gr4ph7]

Those should be the only differences as the implementations are almost entirely shared.

Yeah, seems like it, based on the limited understanding gained from the Rust code.

So it's basically just a tiny addition to the documentation for both functions:

list.agg

Run any polars aggregation expression against the lists’ elements.

This is different from list.eval, which will always wrap even the result of aggregation functions into a list, in that list.agg will return a scalar instead of a single-element list if it is passed an aggregation function.

Behaves like list.eval for non-aggregation expressions.

Examples

If the expression given to .agg ends in an aggregation function, a single value will be returned for each row:

df = pl.DataFrame({"a": [[1, None], [42, 13], [None, None]]})
df.with_columns(null_count=pl.col.a.list.agg(pl.element().null_count()))
# shape: (3, 2)
# ┌──────────────┬────────────┐
# │ a            ┆ null_count │
# │ ---          ┆ ---        │
# │ list[i64]    ┆ u32        │
# ╞══════════════╪════════════╡
# │ [1, null]    ┆ 1          │
# │ [42, 13]     ┆ 0          │
# │ [null, null] ┆ 2          │
# └──────────────┴────────────┘

Otherwise, the result is a list:

df.with_columns(no_nulls=pl.col.a.list.agg(pl.element().drop_nulls()))
# shape: (3, 2)
# ┌──────────────┬───────────┐
# │ a            ┆ no_nulls  │
# │ ---          ┆ ---       │
# │ list[i64]    ┆ list[i64] │
# ╞══════════════╪═══════════╡
# │ [1, null]    ┆ [1]       │
# │ [42, 13]     ┆ [42, 13]  │
# │ [null, null] ┆ []        │
# └──────────────┴───────────┘

list.eval

Run any polars expression against the lists’ elements.

Always returns a list. See list.agg for a specialized version that returns a scalar instead of a single-element list for aggregation functions.