使用 Polars 中的掩码高效替换或重新创建变量

2024-03-16 16:08:38

使用 Polars 中的掩码替换或重新创建变量

在数据分析和转换中，我们经常需要替换或重新创建变量，例如当我们需要用特定值填充缺失值或根据条件设置新值时。Polars 提供了一个强大的掩码机制，使我们能够高效地执行此类操作。

掩码是布尔值数据框，表示数据的真假状态。它们可以用于标识符合特定条件的行或列。在 Polars 中，掩码可以通过使用条件运算符创建，例如：

mask = pl.when(df["amount"] > 270).otherwise(False)

这将创建一个与输入数据框具有相同形状的掩码，其中值为 True 的元素表示满足条件 df["amount"] > 270，而值为 False 的元素表示不满足该条件。

一旦我们有了掩码，就可以使用它来替换现有变量中的值。Polars 提供了 with_columns 方法，它允许我们根据掩码创建新列或修改现有列。例如，要将 amount 列中大于 270 的值替换为 1000，我们可以使用以下代码：

df = df.with_columns([
    pl.col("amount").if_else(mask, 1000, pl.col("amount"))
])

这将创建一个新的 amount 列，其中大于 270 的值替换为 1000，而其他值保持不变。

如果我们希望根据现有变量创建新变量，也可以使用掩码。where 方法允许我们基于掩码选择特定值。例如，要创建一个新的变量 masked_amount，其中仅当 amount 大于 270 时才包含 amount 的值，我们可以使用以下代码：

df = df.with_columns([
    pl.col("masked_amount").where(mask, pl.lit(None))
])

这将创建一个新的 masked_amount 列，其中大于 270 的值包含相应的 amount 值，而其他值则为 None。

掩码机制还可以用于更高级的应用，例如：

累积和 : 通过使用 cum_sum 方法和一个表示累积条件的掩码，我们可以计算累积值。
找出符合特定条件的第一个或最后一个元素 : 通过使用 first_val 和 last_val 方法，我们可以查找符合特定条件的第一个或最后一个元素。
按条件过滤数据框 : 通过使用掩码作为参数传递给 filter 方法，我们可以根据条件过滤数据框。