How to fill a column with random values in polars

I would like to know how to fill a column of a polars dataframe with random values.
The idea is that I have a dataframe with a given number of columns, and I want to add a column to this dataframe which is filled with different random values (obtained from a random.random() function for example).

This is what I tried for now:

df = df.with_columns(pl.when((pl.col('Q') > 0)).then(random.random()).otherwise(pl.lit(1)).alias('Prob'))

With this method, the result that I obtain is a column filled with one random value i.e. all the rows have the same value.

Is there a way to fill the column with different random values ?

Thanks by advance.

You need a "column" of random numbers the same height as your dataframe.

np.random.rand is useful for this:

>>> df.with_columns(random = pl.lit(np.random.rand(df.height)))
shape: (3, 2)
┌─────┬──────────┐
│ foo ┆ random   │
│ --- ┆ ---      │
│ i64 ┆ f64      │
╞═════╪══════════╡
│ 1   ┆ 0.51566  │
│ 2   ┆ 0.009299 │
│ 3   ┆ 0.519169 │
└─────┴──────────┘

>>> df.with_columns(random = pl.when(pl.col("foo") > 2).then(pl.lit(np.random.rand(df.height))))
shape: (3, 2)
┌─────┬──────────┐
│ foo ┆ random   │
│ --- ┆ ---      │
│ i64 ┆ f64      │
╞═════╪══════════╡
│ 1   ┆ null     │
│ 2   ┆ null     │
│ 3   ┆ 0.926295 │
└─────┴──────────┘

First get the number of rows of your dataframe:

row_n = df.select(pl.count()).collect().items()

Then create a random list of that size using random:

to_add = random.sample(range(0, 10), row_n)

And finally add it to your dataframe:

df.with_column(pl.Series(name="new_col", values=to_add))

Create a sample polars dataframe

df = pl.DataFrame({
    'Q': [1, -1, -3, 4],
})

One liner vectorised calculation

df = df.with_columns(
    pl.when(pl.col('Q') > 0)
    .then(pl.lit(np.random.uniform(0, 1, len(df))))
    .otherwise(1)
    .alias('Prob')
)

Result

Q	Prob
1	0.922802
-1	1.0
-3	1.0
4	0.182397

Firstly, you're on a relatively old version of polars if you're still using with_column rather than with_columns so I'd recommend upgrading as there have been new features and performance enhancements. There are also breaking changes, like no more with_column as it was redundant given that it was just a limiting version of with_columns from the beginning.

Setting that aside, and to your issue, the reason it isn't working is that when you run

df.with_columns(pl.when((pl.col('Q') > 0)).then(random.random()).otherwise(pl.lit(1)).alias('Prob'))

python is only calling random.random() once and since it only returns one value, polars broadcasts (ie. copies) it to all the rows. What you need to do is tell python to run it all the times you actually "need". I put need in quotes because polars will complain if you try to give it fewer values than the whole height of the df even though you only need as many random values as there are Q>0.

The easiest way to do that is just with a list comprehension plugging in the height of the df

df.with_columns(
    pl.when((pl.col('Q') > 0))
        .then(pl.lit([random.random() for _ in range(df.height)]))
        .otherwise(pl.lit(1))
        .alias('Prob'))

Using a list comprehension for random.random() isn't as efficient as having numpy create an array of random numbers as it uses optimized C code to do it whereas a list comprehension is just a python loop. I was going for answering the overall question of why it wasn't working rather than to prescribe the method for quickest random number generation.

You could do:

df.with_columns(
    pl.Series(
        [random.random() if q > 0 else 1 for q in df["Q"]]
    ).alias("Prob")
)