Search code examples
pythonpython-polars

Fill several polars columns with a constant value


I am working with the following code...

import polars as pl


df = pl.DataFrame({
    'region': ['GB', 'FR', 'US'],
    'qty':    [3, 6, -8],
    'price':  [100, 102, 95],
    'tenor':   ['1Y', '6M', '2Y'],
})

cols_to_set = ['price', 'tenor']
fill_val    = '-'

df.with_columns([pl.lit(fill_val).alias(c) for c in cols_to_set])

...with the following output.

shape: (3, 4)
┌────────┬─────┬───────┬───────┐
│ region ┆ qty ┆ price ┆ tenor │
│ ---    ┆ --- ┆ ---   ┆ ---   │
│ str    ┆ i64 ┆ str   ┆ str   │
╞════════╪═════╪═══════╪═══════╡
│ GB     ┆ 3   ┆ -     ┆ -     │
│ FR     ┆ 6   ┆ -     ┆ -     │
│ US     ┆ -8  ┆ -     ┆ -     │
└────────┴─────┴───────┴───────┘

Instead of using a list of pl.lit expressions, I thought I could use a single pl.lit(fill_val).alias(cols_to_set). However, this crashes with an error

TypeError: argument 'name': 'list' object cannot be converted to 'PyString'

Is there a way to simplify the above and set all columns in cols_to_set to a specific constant value fill_val?


Solution

  • you could use replace_strict():

    df.with_columns(pl.col(cols_to_set).replace_strict(None, None, default = '-'))
    
    ┌────────┬─────┬───────┬───────┐
    │ region ┆ qty ┆ price ┆ tenor │
    │ ---    ┆ --- ┆ ---   ┆ ---   │
    │ str    ┆ i64 ┆ str   ┆ str   │
    ╞════════╪═════╪═══════╪═══════╡
    │ GB     ┆ 3   ┆ -     ┆ -     │
    │ FR     ┆ 6   ┆ -     ┆ -     │
    │ US     ┆ -8  ┆ -     ┆ -     │
    └────────┴─────┴───────┴───────┘