Polars-Python：比较列表列

Question

在 Python Polars 中，我有一个如下所示的数据框：

df = pl.DataFrame(
    {"sets": [[1, 2, 3], [1, 2], [9, 10]], "optional_members": [[1, 2, 3], [1, 2], [9, 0]]}
)

shape: (2, 2)
┌───────────┬──────────────────┐
│ sets      ┆ optional_members │
│ ---       ┆ ---              │
│ list[i64] ┆ list[i64]        │
╞═══════════╪══════════════════╡
│ [1, 4, 3] ┆ [1, 2, 3]        │
│ [1, 0]    ┆ [1, 2]           │
└───────────┴──────────────────┘

我想构建一个表达式，获取第二列中第一列的元素，保持前一列的形状，即：

shape: (2, 3)
┌───────────┬──────────────────┬─────────────────────┐
│ sets      ┆ optional_members ┆ result              │
│ ---       ┆ ---              ┆ ---                 │
│ list[i64] ┆ list[i64]        ┆ list[bool]          │
╞═══════════╪══════════════════╪═════════════════════╡
│ [1, 4, 3] ┆ [1, 2, 3]        ┆ [true, false, true] │
│ [1, 0]    ┆ [1, 2]           ┆ [true, false]       │
└───────────┴──────────────────┴─────────────────────┘

我尝试在第一个列表上使用 eval，例如：

func = lambda x, y: y.list.contains(x)

df.with_columns(contains=
                pl.col("optional_members")
                .list.
                eval(func(pl.element(), pl.col("optional_members"))))

但是

pl.col()

表达式不能出现在 eval 中。

我们如何解决这个问题？

Answer 1

一种可能是

explode

，比较，然后

group_by.agg

：

(df.with_row_index() # to be able to group back to original rows
   .explode(['sets', 'optional_members'])
   .with_columns(pl.col('sets').eq(pl.col('optional_members')).alias('result'))
   .group_by('index').agg(pl.col(['sets', 'optional_members', 'result']))
   .drop('index')
)

输出：

┌───────────┬──────────────────┬─────────────────────┐
│ sets      ┆ optional_members ┆ result              │
│ ---       ┆ ---              ┆ ---                 │
│ list[i64] ┆ list[i64]        ┆ list[bool]          │
╞═══════════╪══════════════════╪═════════════════════╡
│ [1, 4, 3] ┆ [1, 2, 3]        ┆ [true, false, true] │
│ [1, 0]    ┆ [1, 2]           ┆ [true, false]       │
│ [9, 10]   ┆ [9, 0]           ┆ [true, false]       │
└───────────┴──────────────────┴─────────────────────┘