在“when-then-otherwise”表达式中展开多列

Question

我想用以下逻辑编写一个极坐标的“when-then-otherwise”表达式：

IF a1 and a2 are not null AND b1 and b2 are null 
THEN b1=a1, b2=a2 
ELSE keep values of b1 and b2

作为示例，我有这个数据框

df = pl.DataFrame({
    'a1': [1,1,1], 
    'a2': [2,2,None], 
    'b1': [3,None,None], 
    'b2': [4,None,None]
    })

┌─────┬──────┬──────┬──────┐
│ a1  ┆ a2   ┆ b1   ┆ b2   │
│ --- ┆ ---  ┆ ---  ┆ ---  │
│ i64 ┆ i64  ┆ i64  ┆ i64  │
╞═════╪══════╪══════╪══════╡
│ 1   ┆ 2    ┆ 3    ┆ 4    │
│ 1   ┆ 2    ┆ null ┆ null │
│ 1   ┆ null ┆ null ┆ null │
└─────┴──────┴──────┴──────┘

我想要这个结果

┌─────┬──────┬──────┬──────┐
│ a1  ┆ a2   ┆ b1   ┆ b2   │
│ --- ┆ ---  ┆ ---  ┆ ---  │
│ i64 ┆ i64  ┆ i64  ┆ i64  │
╞═════╪══════╪══════╪══════╡
│ 1   ┆ 2    ┆ 3    ┆ 4    │
│ 1   ┆ 2    ┆ 1    ┆ 2    │
│ 1   ┆ null ┆ null ┆ null │
└─────┴──────┴──────┴──────┘

我尝试了以下方法

df.with_columns(pl.when(pl.col('a1').is_not_null() &
                        pl.col('a2').is_not_null() &
                        pl.col('b1').is_null() &
                        pl.col('b2').is_null()
                       )\
                .then(pl.col('a1','a2'))\
                .otherwise(pl.col('b1','b2'))\
                .name.map(lambda x: {'a1':'b1','a2':'b2'}.get(x))
               )

但这会引发错误

ComputeError: 'with_columns' failed
The reason: expanding more than one `col` is not allowed:

Answer 1

then/otherwise 中不能有多列。

您要么必须对输出的每一列重复

when/then/otherwise

，要么您可以将它们包装在一个结构中，删除旧列，然后取消结构的嵌套。

(
    df
    .with_columns(pl.when(pl.col('a1').is_not_null() &
                        pl.col('a2').is_not_null() &
                        pl.col('b1').is_null() &
                        pl.col('b2').is_null()
                       )
                .then(pl.struct(b1='a1',b2='a2').alias('z'))
                .otherwise(pl.struct('b1','b2').alias('z'))
                )
    .drop('b1','b2')
    .unnest('z')
    )
shape: (3, 4)
┌─────┬──────┬──────┬──────┐
│ a1  ┆ a2   ┆ b1   ┆ b2   │
│ --- ┆ ---  ┆ ---  ┆ ---  │
│ i64 ┆ i64  ┆ i64  ┆ i64  │
╞═════╪══════╪══════╪══════╡
│ 1   ┆ 2    ┆ 3    ┆ 4    │
│ 1   ┆ 2    ┆ 1    ┆ 2    │
│ 1   ┆ null ┆ null ┆ null │
└─────┴──────┴──────┴──────┘