Polars：n_unique()，但作为窗口函数

我需要一种方法来找出特定上下文中两列中有多少个唯一的值对。 基本上像 n_unique，但作为一个窗口函数。

用一个玩具示例来说明：

import polars as pl

dataframe = pl.DataFrame({
    'context': [1, 1, 1,  2, 2, 2,  3, 3, 3],
    'column1': [1, 1, 0,  1, 0, 0,  1, 0, 1],
    'column2': [1, 0, 0,  0, 1, 1,  1, 0, 1]
    # unique:   1  2  3   1  2  -   1  2  -
    # n_unique: -- 3 --   -- 2 --   -- 2 --
})

我想写：

dataframe = (
    dataframe
    .with_columns(
        pl.n_unique('column1', 'column2').over('context').alias('n_unique')
    )
)

从列 ' 的窗口内的列 1、列 2 获取唯一值对的数量语境'。但这不起作用。

我所做的一个尝试是：

(dataframe
    .with_columns(
        pl.concat_list('column1', 'column2').alias('pair')
    )
    .with_columns(
        pl.n_unique('pair').over('context')
    )
)

哪个有效，但是有更好的方法吗？

当然，可以使用 pl.struct 和 pl.n_unique 在 Polars 中实现这一点。与串联列列表不同，使用 pl.struct 应该具有更好的性能，因为它避免了创建中间列表。

以下是如何完成此操作的示例：

dataframe = (
    dataframe
    .with_columns(
        pl.struct(['column1', 'column2']).n_unique().over('context').alias('n_unique')
    )
)
print(dataframe)

这将生成以下输出：

shape: (9, 4)
┌────────┬────────┬────────┬─────────┐
│ context ┆ column1 ┆ column2 ┆ n_unique │
│ ---     ┆ ---     ┆ ---     ┆ ---      │
│ i64     ┆ i64     ┆ i64     ┆ u32      │
╞════════╪════════╪════════╪═════════╡
│ 1       ┆ 1       ┆ 1       ┆ 3        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 1       ┆ 1       ┆ 0       ┆ 3        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 1       ┆ 0       ┆ 0       ┆ 3        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 2       ┆ 1       ┆ 0       ┆ 2        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 2       ┆ 0       ┆ 1       ┆ 2        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 2       ┆ 0       ┆ 1       ┆ 2        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 3       ┆ 1       ┆ 1       ┆ 2        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 3       ┆ 0       ┆ 0       ┆ 2        │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 3       ┆ 1       ┆ 1       ┆ 2        │
└────────┴────────┴────────┴─────────┘

在这种方法中，我们首先使用 pl.struct 从 column1 和 column2 创建一个结构。这使我们能够将两列视为单个实体。然后，我们只需使用 n_unique 函数获取结构的唯一计数，并使用 over 语句将其应用于 context 列定义的每个窗口。

这种方法既高效又可读，使其成为在 Polars 中查找窗口内唯一值对数量的绝佳选择。