Rust polars : unexpected befaviour of when().then().otherwise() in groupby-agg context

Question

I have a complicated mapping logic which I seek to execute within groupby context. The code compiles and doesn't panic, but results are incorrect. I know the logic implementation is correct. Hence, I wonder if when-then-otherwise is supposed to be used within groupby at all?

use polars::prelude::*;
use polars::df;

fn main() {
    let df = df! [
        "Region" => ["EU", "EU", "EU", "EU", "EU"],
        "MonthCCY" => ["APRUSD", "MAYUSD", "JUNEUR", "JULUSD", "APRUSD"],
        "values" => [1, 2, 3, 4, 5],
    ].unwrap();

    let df = df.lazy()
        .groupby_stable([col("MonthCCY")])
        .agg( [
            month_weight().alias("Weight"),
        ]
        );
}

pub fn month_weight() -> Expr {
    when(col("Region").eq(lit("EU")))
    .then(
        // First, If MonthCCY is JUNEUR or AUGEUR - apply 0.05
        when(col("MonthCCY").map( |s|{
            Ok( s.utf8()?
            .contains("JUNEUR|AUGEUR")?
            .into_series() )
         }
            , GetOutput::from_type(DataType::Boolean)
        ))
        .then(lit::<f64>(0.05))
        .otherwise(
            // Second, If MonthCCY is JANEUR - apply 0.0225
            when(col("MonthCCY").map( |s|{
                Ok( s.utf8()?
                .contains("JANEUR")?
                .into_series() )
             }
                , GetOutput::from_type(DataType::Boolean)
            ))
            .then(lit::<f64>(0.0225))
            .otherwise(
                // Third, If MonthCCY starts with JUL or FEB (eg FEBUSD or FEBEUR)- apply 0.15
                when(col("MonthCCY").apply( |s|{
                    let x = s.utf8()?
                    .str_slice(0, Some(3))?;
                    let y = x.contains("JUL|FEB")?
                    .into_series();
                    Ok(y)
                 }
                    , GetOutput::from_type(DataType::Boolean)
                ))
                .then(lit::<f64>(0.15))
                //Finally, if none of the above matched, apply 0.2
                .otherwise(lit::<f64>(0.20))
            )
        )
    ).otherwise(lit::<f64>(0.0))
}

The result I am getting is:

┌──────────┬─────────────┐
│ MonthCCY ┆ Weight      │
│ ---      ┆ ---         │
│ str      ┆ list [f64]  │
╞══════════╪═════════════╡
│ APRUSD   ┆ [0.2, 0.15] │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ MAYUSD   ┆ [0.2]       │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ JUNEUR   ┆ [0.05]      │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ JULUSD   ┆ [0.2]       │
└──────────┴─────────────┘

Clearly, I would expect JULUSD to be [0.15] and APRUSD to be [0.2, 0.2].

Is my expectation of how when().then().otherwise() works within groupby wrong?

I am on Windows11, rustc 1.60.

Answer 1

Yep, you're doing the groupby and the mapping in the wrong order. month_weight() is not an aggregation expression but a simple mapping expression.

As it is, each group of the DataFrame is getting agged into a series that ultimately derives from order of the data in the original frame.

You first want to create a Weight column with values given by the mapping you specify in month_weight(), and then you want to aggregate this column into a list for each group.

So, what you want is the following:

let df = df
    .lazy()
    .with_column(month_weight().alias("Weight")) // create new column first
    .groupby_stable([col("MonthCCY")]) // then group
    .agg([col("Weight").list()]); // then collect into a list per group

println!("{:?}", df.collect().unwrap());

Output:

shape: (4, 2)
┌──────────┬────────────┐
│ MonthCCY ┆ Weight     │
│ ---      ┆ ---        │
│ str      ┆ list [f64] │
╞══════════╪════════════╡
│ APRUSD   ┆ [0.2, 0.2] │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ MAYUSD   ┆ [0.2]      │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ JUNEUR   ┆ [0.05]     │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ JULUSD   ┆ [0.15]     │
└──────────┴────────────┘

---

Also, as an aside, .when().then() can be chained indefinitely; you don't need to nest them. So just as you can write a chained if ... else if ... else if ... else, you can write col().when().then().when().then() ... .otherwise(), which is a lot simpler than nesting each additional condition.