Поляры: n_unique(), но как оконная функция ⇐ Python

[Anonymous]

Мне нужен способ узнать, сколько уникальных пар значений из двух столбцов находится в определенном контексте.
По сути, это похоже на 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')
)
)

Это работает, но есть ли способ лучше?


Подробнее здесь: https://stackoverflow.com/questions/740 ... w-function