如何使用 with_columns 在 lambda 函数中引用多列

Question

有时多列需要在lambda中应用复杂的逻辑过程。 如果 a 或 b 列包含 null 值，则该示例无法运行。（例如这里处理 null 值，实际流程逻辑可能会很复杂）

df = pl.DataFrame({
    "a": [7, 2, 3, 4],
    "b": ['2', '5', None, '7'],})
df.with_columns((pl.col('a')+pl.col('b')).alias('sum'))

结果：

a	b	总和
i64	str	f64
3	空	空
4	7	11

我的期望：

a	b	总和
i64	str	f64
3	空	3
4	7	11

我的解决方案：

def my_sum(x):
    x0= x[0] if x[0] is not None else 0
    x1= x[1] if float(x[1]) is not None else 0
    return x0+x1
df.hstack(df.apply(lambda x:my_sum(x))).rename({"apply": "sum"})

有没有更优雅的方法来完成这项工作？ 如何像 apply 一样使用 with_columns 在 lambda 函数中引用多列？

Answer 1

您可以使用 cast 方法 将字符串列转换为整数，并使用 fill_null 方法 将空值填充为零。 然后执行所需的操作。

方法如下：

df = df.with_columns(pl.col("b").cast(pl.Int64).fill_null(0))
df = df.with_columns((pl.col("a") + pl.col("b")).alias("sum_ab"))

# Output

shape: (4, 3)
┌─────┬─────┬────────┐
│ a   ┆ b   ┆ sum_ab │
│ --- ┆ --- ┆ ---    │
│ i64 ┆ i64 ┆ i64    │
╞═════╪═════╪════════╡
│ 7   ┆ 2   ┆ 9      │
│ 2   ┆ 5   ┆ 7      │
│ 3   ┆ 0   ┆ 3      │
│ 4   ┆ 7   ┆ 11     │
└─────┴─────┴────────┘