when / then / otherwise with values from numpy array - python-polars

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_column(
pl.when(pl.col('group')==3)
.then(values)
.otherwise(pl.col('group')
).alias('group')
)
In [4]: df.with_column(pl.when(pl.col('group')==3).then(values).otherwise(pl.col('group')).alias('group'))
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Cell In [4], line 1
----> 1 df.with_column(pl.when(pl.col('group')==3).then(values).otherwise(pl.col('group')).alias('group'))
File ~/tmp/.venv/lib/python3.8/site-packages/polars/internals/whenthen.py:132, in When.then(self, expr)
111 def then(
112 self,
113 expr: (
(...)
121 ),
122 ) -> WhenThen:
123 """
124 Values to return in case of the predicate being `True`.
125
(...)
130
131 """
--> 132 expr = pli.expr_to_lit_or_expr(expr)
133 pywhenthen = self._pywhen.then(expr._pyexpr)
134 return WhenThen(pywhenthen)
File ~/tmp/.venv/lib/python3.8/site-packages/polars/internals/expr/expr.py:118, in expr_to_lit_or_expr(expr, str_to_lit)
116 return expr.otherwise(None)
117 else:
--> 118 raise ValueError(
119 f"did not expect value {expr} of type {type(expr)}, maybe disambiguate with"
120 " pl.lit or pl.col"
121 )
ValueError: did not expect value [9 8 7] of type <class 'numpy.ndarray'>, maybe disambiguate with pl.lit or pl.col
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_column(
df["group"].set_at_idx(idx, pl.Series(values))
)

Answer 2

Here's all I could come up with
df.with_column(
pl.col("group")
.cumcount()
.over(pl.col("group"))
.alias("idx")
).apply(
lambda x: values[x[1]] if x[0] == 3 else x[0]
).select(
pl.col("apply").alias("group")
)
Surely there's a simpler way?
In [28]: df.with_column(pl.col('group').cumcount().over(pl.col('group')).alias('idx')).apply(lambda x:
...: values[x[1]] if x[0] == 3 else x[0]).select(pl.col('apply').alias('group'))
Out[28]:
shape: (8, 1)
┌───────┐
│ group │
│ --- │
│ i64 │
╞═══════╡
│ 1 │
├╌╌╌╌╌╌╌┤
│ 1 │
├╌╌╌╌╌╌╌┤
│ 1 │
├╌╌╌╌╌╌╌┤
│ 9 │
├╌╌╌╌╌╌╌┤
│ 8 │
├╌╌╌╌╌╌╌┤
│ 7 │
├╌╌╌╌╌╌╌┤
│ 4 │
├╌╌╌╌╌╌╌┤
│ 4 │
└───────┘