when / then / otherwise with values from numpy array

Question

Say I have

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

I have a numpy array of values, which I'd like to replace 'group' 3 with

values = np.array([9, 8, 7])

Here's what I've tried:

(
    df
    .with_columns(
        pl.when(pl.col('group')==3)
        .then(values)
        .otherwise(pl.col('group'))
    ).alias('group')
)

ShapeError: shapes of `self`, `mask` and `other` are not suitable for `zip_with` operation

How can I do this correctly?

Answer 1

A few things to consider.

• One is that you always should convert your numpy arrays to polars Series as we will use the arrow memory specification underneath and not numpys.

• Second is that when -> then -> otherwise operates on columns that are of equal length. We nudge the API in such a direction that you define a logical statement based of columns in your DataFrame and therefore you should not know the indices (nor the lenght of a group) that you want to replace. This allows for much optimizations because if you do not define indices to replace we can push down a filter before that expression.

Anyway, your specific situation does know the length of the group, so we must use something different. We can first compute the indices where the conditional holds and then modify based on those indices.

df = pl.DataFrame({
    "group": [1, 1, 1, 3, 3, 3, 4, 4]
})

values = np.array([9, 8, 7])

# compute indices of the predicate
idx = df.select(
    pl.arg_where(pl.col("group") == 3)
).to_series()

# mutate on those locations
df.with_columns(
    df["group"].scatter(idx, pl.Series(values))
)