pythonpython-polars

Multiply within a group in polars


I would like to take the product within a group in polars

The following works but I suspect there is a more elegant/efficient way to perform this operation. Thank you

import polars as pl
import numpy as np

D = pl.DataFrame({'g':['a','a','b','b'],'v':[1,2,3,4],'v2':[2,3,4,5]})
D.group_by('g').agg(pl.all().map_elements(
                        lambda group: np.prod(group.to_numpy()),return_dtype=pl.Float64))

Solution

  • You could use Expr.product:

    D.group_by('g').agg(pl.all().product())
    

    Output:

    ┌─────┬─────┬─────┐
    │ g   ┆ v   ┆ v2  │
    │ --- ┆ --- ┆ --- │
    │ str ┆ i64 ┆ i64 │
    ╞═════╪═════╪═════╡
    │ b   ┆ 12  ┆ 20  │
    │ a   ┆ 2   ┆ 6   │
    └─────┴─────┴─────┘
    

    If you want Floats:

    D.group_by('g').agg(pl.all().product().cast(pl.Float64))
    
    ┌─────┬──────┬──────┐
    │ g   ┆ v    ┆ v2   │
    │ --- ┆ ---  ┆ ---  │
    │ str ┆ f64  ┆ f64  │
    ╞═════╪══════╪══════╡
    │ b   ┆ 12.0 ┆ 20.0 │
    │ a   ┆ 2.0  ┆ 6.0  │
    └─────┴──────┴──────┘