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       │
└─────────┴─────────┴─────────┴─────────┴─────────┴─────────┘