Группируйте строки DataFrame на основе условий сопоставления нескольких столбцов

Ваш вопрос не совсем ясен, поэтому делаю некоторые предположения.

Во-первых, я предполагаю, что ваше условие таково: группировка должна происходить в пределах col1, а строки принадлежат одной группе, если у них есть одно из совпадений col2, col3.

Во-вторых, неясно, хотите ли вы группировать только последовательные строки или группировать их независимо от их положения.

Итак, вот вам 2 варианта:

Группировка строк независимо от их положения.

Чтобы проиллюстрировать это, DataFrame выглядит следующим образом:

┌──────┬──────┬────────┐
│ col1 ┆ col2 ┆ col3   │
│ ---  ┆ ---  ┆ ---    │
│ i64  ┆ str  ┆ str    │
╞══════╪══════╪════════╡
│ 1    ┆ bra  ┆ brazil │
│ 1    ┆ br   ┆ brasil │
│ 1    ┆ bra  ┆ brazil │
│ 1    ┆ br   ┆ brazil │
└──────┴──────┴────────┘

Будет group_idx следующим образом:

┌──────┬──────┬────────┬───────────┐
│ col1 ┆ col2 ┆ col3   ┆ group_idx │
│ ---  ┆ ---  ┆ ---    ┆ ---       │
│ i64  ┆ str  ┆ str    ┆ i64       │
╞══════╪══════╪════════╪═══════════╡
│ 1    ┆ bra  ┆ brazil ┆ 1         │
│ 1    ┆ br   ┆ brasil ┆ 1         │
│ 1    ┆ bra  ┆ brazil ┆ 1         │
│ 1    ┆ br   ┆ brazil ┆ 1         │
└──────┴──────┴────────┴───────────┘

По сути, вы можете представить строки таблицы как узлы графа, и узлы имеют ребро, связывающее их, если они имеют одинаковый col1 и один из col2, col3 общего. Мы хотим найти все связанные подграфы (см. также Алгоритм заливки).

Для этого вы можете самостоятельно присоединиться к DataFrame на col1 + col2 и col1 + col3 и назначить group_idx и делать это до тех пор, пока вы не объедините все строки, которые имеют что-то общее через другие строки:

dfi = df.with_row_index('group_idx')

while True:
    for col in ['col2','col3']:
        dfi = (
            dfi
            .group_by('col1', col)
            .agg(pl.col('group_idx').min())
            .join(dfi, on=['col1', col], suffix=col)
        )

    if dfi.filter(
           pl.min_horizontal('group_idxcol2', 'group_idxcol3') > pl.col('group_idx')
       ).is_empty():
        break

    dfi = dfi.drop('group_idxcol2', 'group_idxcol3')

dfi.select('col1','col2','col3', pl.col('group_idx').rank('dense'))

┌──────┬──────┬───────────┬───────────┐
│ col1 ┆ col2 ┆ col3      ┆ group_idx │
│ ---  ┆ ---  ┆ ---       ┆ ---       │
│ i64  ┆ str  ┆ str       ┆ u32       │
╞══════╪══════╪═══════════╪═══════════╡
│ 1    ┆ br   ┆ brazil    ┆ 1         │
│ 1    ┆ bra  ┆ brazil    ┆ 1         │
│ 1    ┆ bra  ┆ brasil    ┆ 1         │
│ 1    ┆ col  ┆ collombia ┆ 2         │
│ 1    ┆ col  ┆ columbia  ┆ 2         │
│ 2    ┆ br   ┆ brazil    ┆ 3         │
│ 2    ┆ b    ┆ brazil    ┆ 3         │
└──────┴──────┴───────────┴───────────┘

Группирование только последовательных строк.

Чтобы еще раз проиллюстрировать это, DataFrame выглядит следующим образом:

┌──────┬──────┬────────┐
│ col1 ┆ col2 ┆ col3   │
│ ---  ┆ ---  ┆ ---    │
│ i64  ┆ str  ┆ str    │
╞══════╪══════╪════════╡
│ 1    ┆ bra  ┆ brazil │
│ 1    ┆ br   ┆ brasil │
│ 1    ┆ bra  ┆ brazil │
│ 1    ┆ br   ┆ brazil │
└──────┴──────┴────────┘

Будет group_idx следующим образом:

┌──────┬──────┬────────┬───────────┐
│ col1 ┆ col2 ┆ col3   ┆ group_idx │
│ ---  ┆ ---  ┆ ---    ┆ ---       │
│ i64  ┆ str  ┆ str    ┆ i64       │
╞══════╪══════╪════════╪═══════════╡
│ 1    ┆ bra  ┆ brazil ┆ 1         │
│ 1    ┆ br   ┆ brasil ┆ 2         │ <-- new group cause col2, col3 do not match
│ 1    ┆ bra  ┆ brazil ┆ 3         │ <-- new group cause col2, col3 do not match
│ 1    ┆ br   ┆ brazil ┆ 3         │
└──────┴──────┴────────┴───────────┘

Если эти предположения верны, то вы можете сравнить значения строк с предыдущими значениями, используя Expr.shift(), чтобы проверить, выполняется ли условие (col1 остается тем же, один из col2 или col3 все тот же).

Итак, сначала вы можете представить себе следующий DataFrame:

(
    df
    .with_columns(
        col1_change=pl.col('col1') != pl.col('col1').shift(1),
        col2_change=pl.col('col2') != pl.col('col2').shift(1),
        col3_change=pl.col('col3') != pl.col('col3').shift(1),
    ).fill_null(True)
)

┌──────┬──────┬───────────┬─────────────┬─────────────┬─────────────┐
│ col1 ┆ col2 ┆ col3      ┆ col1_change ┆ col2_change ┆ col3_change │
│ ---  ┆ ---  ┆ ---       ┆ ---         ┆ ---         ┆ ---         │
│ i64  ┆ str  ┆ str       ┆ bool        ┆ bool        ┆ bool        │
╞══════╪══════╪═══════════╪═════════════╪═════════════╪═════════════╡
│ 1    ┆ br   ┆ brazil    ┆ true        ┆ true        ┆ true        │
│ 1    ┆ bra  ┆ brazil    ┆ false       ┆ true        ┆ false       │
│ 1    ┆ bra  ┆ brasil    ┆ false       ┆ false       ┆ true        │
│ 1    ┆ col  ┆ collombia ┆ false       ┆ true        ┆ true        │
│ 1    ┆ col  ┆ columbia  ┆ false       ┆ false       ┆ true        │
│ 2    ┆ br   ┆ brazil    ┆ true        ┆ true        ┆ true        │
│ 2    ┆ b    ┆ brazil    ┆ false       ┆ true        ┆ false       │
└──────┴──────┴───────────┴─────────────┴─────────────┴─────────────┘

Теперь нас не волнует, изменится ли col2, а col3 нет, и наоборот. Чтобы настроить это условие, мы можем использовать .min_horizontal():

(
    df
    .with_columns(
        col1_change=pl.col('col1') != pl.col('col1').shift(1),
        col2_or_col3_change=pl.min_horizontal(pl.col('col2','col3') != pl.col('col2','col3').shift(1))
    ).fill_null(True)
)

┌──────┬──────┬───────────┬─────────────┬─────────────────────┐
│ col1 ┆ col2 ┆ col3      ┆ col1_change ┆ col2_or_col3_change │
│ ---  ┆ ---  ┆ ---       ┆ ---         ┆ ---                 │
│ i64  ┆ str  ┆ str       ┆ bool        ┆ bool                │
╞══════╪══════╪═══════════╪═════════════╪═════════════════════╡
│ 1    ┆ br   ┆ brazil    ┆ true        ┆ true                │
│ 1    ┆ bra  ┆ brazil    ┆ false       ┆ false               │
│ 1    ┆ bra  ┆ brasil    ┆ false       ┆ false               │
│ 1    ┆ col  ┆ collombia ┆ false       ┆ true                │
│ 1    ┆ col  ┆ columbia  ┆ false       ┆ false               │
│ 2    ┆ br   ┆ brazil    ┆ true        ┆ true                │
│ 2    ┆ b    ┆ brazil    ┆ false       ┆ false               │
└──────┴──────┴───────────┴─────────────┴─────────────────────┘

Теперь нам нужно создать новую группу, когда одно из этих условий истинно, и для этого мы можем использовать либо |, либо .max_horizontal(). Я использую .max_horizontal(), потому что он позволяет мне переключаться с DataFramw.fill_null() на Expr.fill_null() и это пригодится на следующем этапе:

(
    df
    .with_columns(
        group_start = pl.max_horizontal(
            (pl.col('col1') != pl.col('col1').shift(1)) |
            pl.min_horizontal(pl.col('col2','col3') != pl.col('col2','col3').shift(1))
        ).fill_null(True)
    )
)

┌──────┬──────┬───────────┬─────────────┐
│ col1 ┆ col2 ┆ col3      ┆ group_start │
│ ---  ┆ ---  ┆ ---       ┆ ---         │
│ i64  ┆ str  ┆ str       ┆ bool        │
╞══════╪══════╪═══════════╪═════════════╡
│ 1    ┆ br   ┆ brazil    ┆ true        │
│ 1    ┆ bra  ┆ brazil    ┆ false       │
│ 1    ┆ bra  ┆ brasil    ┆ false       │
│ 1    ┆ col  ┆ collombia ┆ true        │
│ 1    ┆ col  ┆ columbia  ┆ false       │
│ 2    ┆ br   ┆ brazil    ┆ true        │
│ 2    ┆ b    ┆ brazil    ┆ false       │
└──────┴──────┴───────────┴─────────────┘

И теперь вы можете использовать Expr.cum_sum() для перечисления групп:

(
    df
    .with_columns(
        group_idx = pl.max_horizontal(
            (pl.col('col1') != pl.col('col1').shift(1)) |
            pl.min_horizontal(pl.col('col2','col3') != pl.col('col2','col3').shift(1))
        ).fill_null(True).cum_sum()
    )
)

┌──────┬──────┬───────────┬───────────┐
│ col1 ┆ col2 ┆ col3      ┆ group_idx │
│ ---  ┆ ---  ┆ ---       ┆ ---       │
│ i64  ┆ str  ┆ str       ┆ u32       │
╞══════╪══════╪═══════════╪═══════════╡
│ 1    ┆ br   ┆ brazil    ┆ 1         │
│ 1    ┆ bra  ┆ brazil    ┆ 1         │
│ 1    ┆ bra  ┆ brasil    ┆ 1         │
│ 1    ┆ col  ┆ collombia ┆ 2         │
│ 1    ┆ col  ┆ columbia  ┆ 2         │
│ 2    ┆ br   ┆ brazil    ┆ 3         │
│ 2    ┆ b    ┆ brazil    ┆ 3         │
└──────┴──────┴───────────┴───────────┘