python-polars - 如何动态创建字段

Question

有什么方法可以动态创建字段吗？我知道有一些方法。但最好了解极地的最佳方法。例如，我想向现有数据帧添加 12 个移位列。(lag1, lag2, lag3...lagN) 如何实现这一点？

谢谢。

Answer 1

您可以使用 python 语言来实现这一点。 Polars 表达式是延迟求值的，因此您可以在任何地方创建它们，在 for 循环、函数、列表理解中等等。

下面我给出了一个动态创建lag列的示例，其中一个方法是调用函数，分配给变量，然后使用该变量。并且具有列表理解能力。

# some initial dataframe
df = pl.DataFrame({
    "a": [1, 2, 3, 4, 5],
    "b": [5, 4, 3, 2, 1]
})

# a function that returns a lazy evaluated expression
def lag(name: str, n: int) -> pl.Expr:
    return pl.col(name).shift(n).suffix(f"_lag_{n}")

# a lazy evaluated expression assigned to a variable
lag_foo = lag("a", 1)

out = df.select([
    lag_foo,
] + [lag("b", i) for i in range(5)]  # create exprs with a list comprehension
)

print(out)

输出:

shape: (5, 6)
┌─────────┬─────────┬─────────┬─────────┬─────────┬─────────┐
│ a_lag_1 ┆ b_lag_0 ┆ b_lag_1 ┆ b_lag_2 ┆ b_lag_3 ┆ b_lag_4 │
│ ---     ┆ ---     ┆ ---     ┆ ---     ┆ ---     ┆ ---     │
│ i64     ┆ i64     ┆ i64     ┆ i64     ┆ i64     ┆ i64     │
╞═════════╪═════════╪═════════╪═════════╪═════════╪═════════╡
│ null    ┆ 5       ┆ null    ┆ null    ┆ null    ┆ null    │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 1       ┆ 4       ┆ 5       ┆ null    ┆ null    ┆ null    │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 2       ┆ 3       ┆ 4       ┆ 5       ┆ null    ┆ null    │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 3       ┆ 2       ┆ 3       ┆ 4       ┆ 5       ┆ null    │
├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 4       ┆ 1       ┆ 2       ┆ 3       ┆ 4       ┆ 5       │
└─────────┴─────────┴─────────┴─────────┴─────────┴─────────┘