[Python] Does head(n) in Polars agg guarantee alignment across multiple columns within the...

When using group_by().agg() in Polars to apply head(n) on multiple columns simultaneously, is it guaranteed that the returned elements come from the same original rows?

For example, when grouping by code and applying .head(2) to several columns (e.g., 'interest', 'embedding', 'a'), will the result contain aligned values — such as i1,e1,a1 and i2,e2,a2 — corresponding to the first two rows within each group?

I’ve tested this with sample code below, and it behaves as expected. But I’m wondering: is this alignment behavior officially guaranteed, or is it an implementation detail that could change?

import polars as pl

# 示例数据
df = pl.DataFrame({
'code': ['A', 'A', 'A', 'B', 'B'],
'interest': ['i1', 'i2', 'i3', 'i4', 'i5'],
'embedding': ['e1', 'e2', 'e3', 'e4', 'e5'],
'a': ['a1', 'a2', 'a3', 'a4', 'a5']
})

# 分组并获取前2个样本
result = (
df.group_by('code')
.agg([
pl.col('interest').head(2).alias('interests'),
pl.col('embedding').head(2).alias('embeddings'),
pl.col('a').head(2).alias('a')
])
)

print(result)

Continue reading...