polars assign to multiple columns

Question

I am trying to understand if there is any way to do when..then..otherwise in polars and assign to multiple columns. I have a elo dataset with millions of rows where I want to assign the current elo to anything greater than date. In pandas, I would do

elo_df.loc[(id, date:)), ["elo", "true_skill_mu", "true_skill_sigma"]] = elo, true_skill_mu, true_skill_sigma

The code below works but is very slow. I am hoping I can increase the speed by at least 3x by making the filter happen once. Also, if have any suggestion on to how tomake this faster, please let me know.

elo_df = elo_df.with_columns([pl.when((pl.col("id") == col) & (pl.col("date") >= date)).then(pl.lit(new_rating)).otherwise(pl.col("elo")).alias("elo"),
                                           
pl.when((pl.col("id") == col) & (pl.col("date") >= date)).then(pl.lit(new_mu)).otherwise(pl.col("true_skill_mu")).alias("true_skill_mu"),
                                            pl.when((pl.col("id") == col) & (pl.col("date") >= date)).then(pl.lit(new_sigma)).otherwise(pl.col("true_skill_sigma")).alias("true_skill_sigma")]

Answer 1

use polars.col followed by polars.Expr.map

def on_when(s : pl.Series, df = elo_df):
    # or pass them as default arguments
    global new_rating, new_mu, new_sigma
    if s.name == "elo":
        return new_rating
    elif s.name == "true_skill_mu":
        return new_mu
    elif s.name == "true_skill_sigma":
        return new_sigma


def otherwise(s : pl.Series, df = elo_df):
    if s.name == "elo":
        return df.get_column("elo")
    elif s.name == "true_skill_mu":
        return df.get_column("true_skill_mu")
    elif s.name == "true_skill_sigma":
        return df.get_column("true_skill_sigma")
    # or you can simply just use `return s` since we are not doing any operations
    # this is just an example if we want to do calculations based on other columns


elo_df = elo_df.with_columns([
    pl.when((pl.col("id") == col) & (pl.col("date") >= date))
    .then(pl.col(["elo", "true_skill_mu", "true_skill_sigma"]).map(on_when))
    .otherwise(pl.col(["elo", "true_skill_mu", "true_skill_sigma"]).map(otherwise))
])