Polars：使用列值在when / then 表达式中引用其他列

Question

我有一个 Polars 数据框，我想在其中使用“when/then”表达式派生一个新列。新列的值应取自同一数据帧中的不同列。但是，从行中获取值的列因行而异。

这是一个简单的例子：

df = pl.DataFrame(
    {
        "frequency": [1, None, None, None],
        "frequency_ref": ["a", "b", "a", "a"], # this column signals from which other column values should be taken in case there's no value in the "frequency" column
        "a": [1, 2, 3, 4]
        "b": [5, 6, 7, 8],
    }
)

生成的数据框应如下所示：

res = pl.DataFrame(
    {
        "frequency": [0.5, None, None, None],
        "frequency_ref": ["a", "z", "a", "a"],
        "a": [1, 2, 3, 4],
        "z": [5, 6, 7, 8],
        "res": [0.5, 6, 3, 4]
    }
)

我尝试使用嵌套的 pl.col 创建动态引用：

# Case 1) Fixed value is given
fixed_freq_condition = pl.col("frequency").is_not_null() & pl.col("frequency").is_not_nan()
# Case 2) Reference to distribution data is given
ref_freq_condition = pl.col("frequency_ref").is_not_null()

# Apply the conditions to calculate res
df = df.with_columns(
    pl.when(fixed_freq_condition)
    .then(pl.col("frequency"))
    .when(ref_freq_condition)
    .then(
      pl.col(pl.col("frequency_ref"))
    )
    .otherwise(0.0)
    .alias("res"),
)

失败了

TypeError: invalid input for "col". Expected "str" or "DataType", got 'Expr'.

有效的方法（但仅作为中间解决方案）是在一个很长的when/then 表达式中显式列出每个可能的列值。这远非最佳，因为列名称将来可能会发生变化并产生大量代码重复。

df = df.with_columns(
    pl.when(fixed_freq_condition)
    .then(pl.col("frequency"))
    .when(pl.col("frequency_ref") == "a")
    .then(pl.col("a"))
    # ... more entries
    .when(pl.col("frequency_ref") == "z")
    .then(pl.col("z"))
    .otherwise(0.0)
    .alias("res"),
)

Answer 1

您可以手动列出列，尽管这不太符合人体工程学：

pl.when(fixed_freq_condition)
.then(pl.col("frequency"))
.when(ref_freq_condition)
.then(
    pl.when(pl.col("frequency_ref") == "a")
    .then(pl.col("a"))
    .when(pl.col("frequency_ref") == "b")
    .then(pl.col("b"))
    .. etc etc ..
    .otherwise(0.0)
)
.otherwise(0.0)
.alias("res"),