如何在Polars LazyFrame中按方差>0的条件筛选列？

polars.LazyFrame.var 方法会返回表中每一列的方差值，示例如下：

>>> df = pl.DataFrame({"a": [1, 2, 3, 4], "b": [1, 2, 1, 1], "c": [1, 1, 1, 1]}).lazy()
>>> df.collect()
shape: (4, 3)
┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 1   ┆ 1   │
│ 2   ┆ 2   ┆ 1   │
│ 3   ┆ 1   ┆ 1   │
│ 4   ┆ 1   ┆ 1   │
└─────┴─────┴─────┘
>>> df.var().collect()
shape: (1, 3)
┌──────────┬──────┬─────┐
│ a        ┆ b    ┆ c   │
│ ---      ┆ ---  ┆ --- │
│ f64      ┆ f64  ┆ f64 │
╞══════════╪══════╪═════╡
│ 1.666667 ┆ 0.25 ┆ 0.0 │
└──────────┴──────┴─────┘

想要在LazyFrame中直接筛选出方差大于0的列，在普通DataFrame中可以通过遍历列实现：

>>> data = df.collect()
>>> data.var()
shape: (1, 3)
┌──────────┬──────┬─────┐
│ a        ┆ b    ┆ c   │
│ ---      ┆ ---  ┆ --- │
│ f64      ┆ f64  ┆ f64 │
╞══════════╪══════╪═════╡
│ 1.666667 ┆ 0.25 ┆ 0.0 │
└──────────┴──────┴─────┘
>>> cols = pl.select([s for s in data.var() if (s > 0).all()]).columns
>>> cols
['a', 'b']
>>> data.select(cols)
shape: (4, 2)
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 1   │
│ 2   ┆ 2   │
│ 3   ┆ 1   │
│ 4   ┆ 1   │
└─────┴─────┘

但上述方法直接套用到LazyFrame中会报错，因为LazyFrame对象不支持下标遍历：

>>> data = data.lazy()
>>> data
<LazyFrame at 0x138573D30>
>>> cols = pl.select([s for s in data.var() if (s > 0).all()]).columns
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 1, in <listcomp>
  File "/home/jasmine/miniconda3/envs/jupyternb/lib/python3.9/site-packages/polars/internals/lazyframe/frame.py", line 421, in __getitem__
    raise TypeError(
TypeError: 'LazyFrame' object is not subscriptable (aside from slicing). Use 'select()' or 'filter()' instead.

解决方案 #

要在LazyFrame中高效实现该需求，可以通过以下步骤完成，全程保持懒执行的性能优势：

1. 一次性计算所有列的方差，转置结果后筛选出方差大于0的列名；
2. 使用筛选得到的列名从原LazyFrame中选择目标列。

具体代码如下：

import polars as pl

df = pl.DataFrame({"a": [1, 2, 3, 4], "b": [1, 2, 1, 1], "c": [1, 1, 1, 1]}).lazy()

# 1. 获取方差>0的列名（仅在collect时执行一次方差计算）
valid_cols = (
    df.var()
    .transpose(include_header=True, header_name="column", value_name="variance")
    .filter(pl.col("variance") > 0)
    .select("column")
    .collect()
    .to_series()
    .to_list()
)

# 2. 用筛选后的列名选择原LazyFrame的列
filtered_df = df.select(valid_cols).collect()
print(filtered_df)

执行结果：

shape: (4, 2)
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 1   │
│ 2   ┆ 2   │
│ 3   ┆ 1   │
│ 4   ┆ 1   │
└─────┴─────┘

这种方法只需要计算一次所有列的方差，避免了重复计算，同时保留了LazyFrame的性能优势，适合处理大规模数据集。

内容的提问来源于stack exchange，提问作者JASMINE LIAW