如何在列表类型的 DataFrame 列中将 `[1, 5]` 转换为 `[1, 2, 3, 4, 5]`？

Question

我有一个极地数据框：

pl.DataFrame({'a':[[1,3], [1,5]]})

a
list
[1, 3]
[1, 5]

我想做一些矢量化操作来将其扩展为：

a
list
[1, 2, 3]
[1, 2, 3, 4, 5]

我想出的一个解决方案是将数组分成两列（init 和 final），然后执行 pl.struct(['init', 'final'])，然后执行 apply获取范围。

def get_valid_codes(struct: dict) -> list:
    code_range = set(range(struct['init'], struct['final'] + 1))
    codes      =  list(set.intersection(valid_codes, code_range))
    return codes if codes else [0]

这对我的数据集（3 亿行）来说很慢，我想知道是否有更好的方法。

如果您能弄清楚如何从列表中过滤掉某些（预定义的）值，则加分。

Answer 1

让我们扩展数据，以便我们可以显示 'bad codes' 的一些逻辑。

import polars as pl

df = pl.DataFrame({"a": [[1, 3], [1, 5], [7, 9], [3, 7], [9, 13], [5, 11]]})
print(df)

shape: (6, 1)
┌────────────┐
│ a          │
│ ---        │
│ list [i64] │
╞════════════╡
│ [1, 3]     │
├╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [1, 5]     │
├╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [7, 9]     │
├╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [3, 7]     │
├╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [9, 13]    │
├╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [5, 11]    │
└────────────┘

我们将使用 6 到 10 作为 'bad codes' 进行淘汰。我们将大量使用表达式的 arr 子集：

bad_codes = [6, 7, 8, 9, 10]
df.with_column(
    pl.arange(pl.col("a").arr.first(), pl.col("a").arr.last() + 1)
    .arr.eval(pl.first().filter(pl.first().is_in(bad_codes).is_not()),
              parallel=True)
    .alias("result")
)

shape: (6, 2)
┌────────────┬───────────────┐
│ a          ┆ result        │
│ ---        ┆ ---           │
│ list [i64] ┆ list [i64]    │
╞════════════╪═══════════════╡
│ [1, 3]     ┆ [1, 2, 3]     │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [1, 5]     ┆ [1, 2, ... 5] │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [7, 9]     ┆ []            │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [3, 7]     ┆ [3, 4, 5]     │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [9, 13]    ┆ [11, 12, 13]  │
├╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ [5, 11]    ┆ [5, 11]       │
└────────────┴───────────────┘

当所有代码都是“错误代码”时，此算法会留下一个空列表 []。如果您需要 [0] 而不是空列表，您可以使用 pl.when 和 arr.lengths 表达式将它们更改为 [0].

如何在列表类型的 DataFrame 列中将 `[1, 5]` 转换为 `[1, 2, 3, 4, 5]`？

How to turn `[1, 5]` into `[1, 2, 3, 4, 5]` in a DataFrame column of list type?

python-polars