删除特定列后的所有列

Question

假设我正在读取一个有几百列的csv文件。

现在，我知道在某一列，比如说 'XYZ' 之后，所有的列都是无用的。

我想保留从开始到 'XYZ' 这一列的所有列，并删除 'XYZ' 之后的所有列。

在pandas中，我可以这样做：

df.iloc[:, :df.columns.get_loc('XYZ') + 1]

那么在polars中，有什么高效的方法呢？

Answer 1

如果你确定这个列是存在的，你可以像使用pandas那样进行操作：

df = pl.DataFrame({'ABC': [1,2,3], 'DEF': [4,5,6],
                   'XYZ': [7,8,9], 'GHI': [10,11,12]})

out = df[:, :df.columns.index('XYZ')+1]
# or 
out = df[:, :df.find_idx_by_name('XYZ')+1]

或者，你也可以用更简短（而且更高效）的方法：

out = df[:, :'XYZ']

输出结果：

shape: (3, 3)
┌─────┬─────┬─────┐
│ ABC ┆ DEF ┆ XYZ │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 4   ┆ 7   │
│ 2   ┆ 5   ┆ 8   │
│ 3   ┆ 6   ┆ 9   │
└─────┴─────┴─────┘

Answer 2

用 [:] 这种切片方法在急切求值的情况下是有效的，适用于 DataFrame，但我知道它对 LazyFrame 并不太管用。

如果你想使用懒惰求值，可以直接使用 DataFrame.select() 方法：

# prepare the data
df = pl.LazyFrame({
    'ABC': [1,2,3],
    'DEF': [4,5,6],
    'XYZ': [7,8,9],
    'garbage1': [10,11,12],
    'garbage2': list('abc')
})
df.sink_csv('test.csv')

现在我们可以进行扫描了：

df = pl.scan_csv('test.csv')
df.select(df.columns[:df.columns.index('XYZ')+1]).collect()

┌─────┬─────┬─────┐
│ ABC ┆ DEF ┆ XYZ │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 4   ┆ 7   │
│ 2   ┆ 5   ┆ 8   │
│ 3   ┆ 6   ┆ 9   │
└─────┴─────┴─────┘

不过，我不太确定 polars 在过滤掉那些不需要返回的列时效率如何，因为这些列还是得从文件中读取。