python-polars - Polars 中的表达式选择引用早期别名的上下文

Question

有没有办法允许 Polars 中的表达式引用以前的别名表达式？例如，这段代码定义了两个新列错误，因为第二个新列引用了第一个：

import polars as pl

df = pl.DataFrame(dict(x=[0, 0, 1]))

df.select([
  (pl.col('x') + 1).alias('y'),
  (pl.col('y') * 2).alias('z')],
)
# pyo3_runtime.PanicException: called `Result::unwrap()` on an `Err` value:
# NotFound("Unable to get field named \"y\". Valid fields: [\"x\"]")

该错误表明失败是由第一个别名对第二个表达式不可见引起的。有没有一种简单的方法来完成这项工作？

Answer 1

上下文中的所有极坐标表达式都是并行执行的。所以他们不能引用一个尚不存在的列。

一个上下文是：

• df.with_columns
• df.select
• df.groupby(..).agg

这意味着您需要对引用其他表达式输出的表达式强制执行顺序执行。

在你的情况下，我会这样做：

(df.with_column(
       (pl.col('x') + 1).alias('y')
    ).select([
       pl.col('y'),
      (pl.col('y') * 2).alias('z')
]))

Answer 2

一种解决方法是将每个新列拉出到自己的with_column调用中，然后进行最终select保留以保留您应该保留的列。您可能希望确保这是懒惰地完成的。

import polars as pl

df = pl.DataFrame(dict(x=[0, 0, 1]))

(df
  .lazy()
  .with_column((pl.col("x") + 1).alias("y"))
  .with_column((pl.col("y") * 2).alias("z"))
  .select(["y", "z"])
  .collect()
)
# shape: (3, 2)
# ┌─────┬─────┐
# │ y   ┆ z   │
# │ --- ┆ --- │
# │ i64 ┆ i64 │
# ╞═════╪═════╡
# │ 1   ┆ 2   │
# ├╌╌╌╌╌┼╌╌╌╌╌┤
# │ 1   ┆ 2   │
# ├╌╌╌╌╌┼╌╌╌╌╌┤
# │ 2   ┆ 4   │
# └─────┴─────┘