Polars中when-then-otherwise异常及单表达式实现咨询

问题场景 #

测试代码 #

import polars as pl
import numpy as np

u = 2

df = (pl.DataFrame(dict(
  j=np.random.randint(10, 99, 10)
  ))
  .with_row_count()
  .with_columns(k=(pl.col('j') % 10) == u)
  .with_columns(l=pl.col('k').any())
  )
print(df)

print(df
  .with_columns(i=pl
    .when(pl.col('k').any())
    .then(pl.col('row_nr') > pl.col('row_nr').where(pl.col('k')).first())
    .otherwise(None)
    )
  )

异常情况 #

当设置u=0且j%10无匹配项时，触发如下错误：

exceptions.ComputeError: cannot evaluate two series of different lengths (10 and 0)

Error originated in expression: '[(col("row_nr")) > (col("row_nr").filter(col("k")).first())]'

疑问解答 #

1. 为何pl.when().then().otherwise()在when条件为false时仍会求值then分支？ #

Polars的表达式执行逻辑和Python原生if/else不同：Polars会预先解析并计算所有分支的表达式，再根据when的条件选择最终结果，而非跳过不满足条件的分支。

也就是说，无论when条件是否成立，then和otherwise里的表达式都会先被执行一遍。当u=0且无匹配项时，pl.col('row_nr').where(pl.col('k'))生成长度为0的Series，调用first()后仍为长度0的Series，和长度10的row_nr做比较时就会抛出长度不匹配的错误。

2. 能否在单个表达式内处理该场景，无需外部Python逻辑？ #

可以，核心是确保then分支的表达式在任何情况下都能生成合法的、长度匹配的Series。通过fill_null或coalesce给空Series的first()结果设置默认值，就能避免报错：

print(df
  .with_columns(i=pl
    .when(pl.col('k').any())
    .then(pl.col('row_nr') > pl.col('row_nr').where(pl.col('k')).first().fill_null(-1))
    .otherwise(None)
    )
  )

或者用coalesce更明确地处理空值：

print(df
  .with_columns(i=pl
    .when(pl.col('k').any())
    .then(pl.col('row_nr') > pl.coalesce(pl.col('row_nr').where(pl.col('k')).first(), pl.lit(-1)))
    .otherwise(None)
    )
  )

因为when的条件是pl.col('k').any()，只有存在匹配项时才会进入then分支，此时first()会返回有效值，默认值-1不会被实际使用；无匹配项时，otherwise直接返回None，完全避开了then分支的计算结果。

内容的提问来源于stack exchange，提问作者levant pied