Broadcast in agg when needed

Question

In Polars, the select and with_column methods broadcast any scalars that they get, including literals:

import polars as pl

df.with_column(pl.lit(1).alias("y"))
# shape: (3, 2)
# ┌─────┬─────┐
# │ x   ┆ y   │
# │ --- ┆ --- │
# │ i64 ┆ i64 │
# ╞═════╪═════╡
# │ 1   ┆ 1   │
# ├╌╌╌╌╌┼╌╌╌╌╌┤
# │ 2   ┆ 1   │
# ├╌╌╌╌╌┼╌╌╌╌╌┤
# │ 3   ┆ 1   │
# └─────┴─────┘

The agg method does not broadcast literals:

import polars as pl

df = pl.DataFrame(dict(x=[1,1,0,0])).groupby("x")

df.agg(pl.lit(1).alias("y"))
# exceptions.ComputeError: returned aggregation is a different length: 1 than the group lengths: 2

Is there an operation I can apply that will broadcast a scalar and ignore a non-scalar? Something like this:

df.agg(something(pl.lit(1)).alias("y"))
# shape: (2, 2)
# ┌─────┬─────┐
# │ x   ┆ y   │
# │ --- ┆ --- │
# │ i64 ┆ i64 │
# ╞═════╪═════╡
# │ 0   ┆ 1   │
# ├╌╌╌╌╌┼╌╌╌╌╌┤
# │ 1   ┆ 1   │
# └─────┴─────┘

score 1 · Answer 1 · answered Mar 27 '22 at 08:19

1

You will need to use pl.repeat(1, pl.count()) to expand the literal to the group size.

(answer from the Polars Issue tracker - https://github.com/pola-rs/polars/issues/2987#issuecomment-1079617229)

answered Mar 27 '22 at 08:19

jvz

1,183
6
13

I posted this question mainly because I could not get that answer from the issue tracker to work. It (1) results in a column of lists when run in `agg`, so it can't be used in `agg`, (2) errors unless you unwrap the literal, so it won't work on expressions, and (3) does not ignore expressions that are already the correct size, so it's not a broadcast operator, per se. – drhagen Mar 27 '22 at 11:57

lemmingxuan · Answer 2 · 2022-05-09T02:05:25.237

0

df = pl.DataFrame(dict(x=[1,1,0,0])).groupby("x")
df.agg( pl.repeat(1, pl.count() ) ).explode(
    pl.col('literal')
)

may be helpful

edited May 09 '22 at 02:05

answered May 09 '22 at 01:57

lemmingxuan

549
1
7
18

score 0 · Accepted Answer · answered Jul 05 '23 at 00:37

Upgrade to 0.13.51 or later. The original behavior was considered a bug, and now literals are broadcast.

import polars as pl

df = pl.DataFrame(dict(x=[1,1,0,0])).groupby("x")

df.agg(pl.lit(1).alias("y"))
# shape: (2, 2)
# ┌─────┬─────┐
# │ x   ┆ y   │
# │ --- ┆ --- │
# │ i64 ┆ i32 │
# ╞═════╪═════╡
# │ 0   ┆ 1   │
# │ 1   ┆ 1   │
# └─────┴─────┘

Broadcast in agg when needed

3 Answers3