Векторное обновление статусов в pandas: transform и map против merge

Как в pandas обновлять статус в DataFrame по условиям внутри группы без groupby.apply: векторные решения transform+map и merge, проверка related_id в группе.

Обновление значения в столбце DataFrame по условиям, затрагивающим несколько строк внутри одной группы, — задача распространённая, но коварная, когда нужно ссылаться на другие строки. Ниже — простой и быстрый способ сделать это в pandas без медленных построчных операций.

Постановка задачи

У нас есть данные, сгруппированные по group_id. Для конкретной строки мы хотим присвоить status значение "resolved" только при выполнении всех условий: текущая строка имеет статус pending, в этой же группе есть хотя бы одна строка со статусом active, а related_id ссылается на строку из той же группы. Пример данных ниже.

import pandas as pd
records = {
    'id': [1, 2, 3, 4, 5, 6, 7],
    'group_id': ['A', 'A', 'A', 'B', 'B', 'B', 'B'],
    'status': ['pending', 'active', 'pending', 'pending', 'active', 'pending', 'pending'],
    'related_id': [None, None, 1, None, None, 4, 4]
}
frame = pd.DataFrame(records)
print(frame)

Ожидаемый результат: в строках 3, 6 и 7 статус должен измениться на resolved.

Почему наивные подходы с группами дают сбои

Использовать groupby().apply() заманчиво, но неэффективно: операция медленная и здесь почти никогда не нужна. Вдобавок формулировка условия неоднозначна: «в группе есть хотя бы одна строка со статусом active, а related_id совпадает с id текущей строки». В примере строки, на которые указывает related_id (id 1 и 4), имеют статус pending, а не active. Отсюда два прочтения: либо мы проверяем, что related_id относится к этой же группе, а в группе вообще есть активная строка; либо явно находим строку по related_id и проверяем её статус. Оба варианта можно реализовать эффективно.

Решение 1: векторные проверки с transform и map

Первое прочтение: помечаем pending‑строку как resolved, если её related_id указывает на запись из той же группы и в группе есть хотя бы одна строка со статусом active. Так мы обходимся без apply и используем чисто векторные операции.

# Текущая строка имеет статус pending?
cond_pending = frame['status'].eq('pending')
# related_id указывает на строку в той же группе?
cond_same_group = frame['related_id'].map(frame.set_index('id')['group_id']).eq(frame['group_id'])
# В этой группе есть хотя бы одна строка со статусом active?
cond_active_exists = frame['status'].eq('active').groupby(frame['group_id']).transform('any')
# Обновляем статус, где выполнены все условия
frame.loc[cond_pending & cond_same_group & cond_active_exists, 'status'] = 'resolved'
print(frame)

Ожидаемый вывод:

   id group_id    status  related_id
0   1        A   pending         NaN
1   2        A    active         NaN
2   3        A  resolved         1.0
3   4        B   pending         NaN
4   5        B    active         NaN
5   6        B  resolved         4.0
6   7        B  resolved         4.0

Это соответствует целевому результату и работает быстро благодаря полностью векторной логике.

Решение 2: явный поиск через merge

Если нужно явно получить строку, на которую ссылается related_id, и проверить её статус, используйте merge. В этой версии мы помечаем pending‑строку как resolved, когда строка, на которую она ссылается (по related_id в той же группе), имеет статус pending.

# Маска pending для текущих строк
mask_pending = frame['status'].eq('pending')
# Для pending-строк находим ссылочную строку (в той же группе) и проверяем её статус
mask_lookup = (
    frame.loc[mask_pending, ['related_id', 'group_id']].reset_index()
        .merge(
            frame[['id', 'group_id', 'status']],
            left_on=['related_id', 'group_id'],
            right_on=['id', 'group_id']
        )
        .set_index('index')['status']
        .eq('pending')
)
# Применяем обновление для строк, где выполнены оба условия
frame.loc[mask_pending & mask_lookup, 'status'] = 'resolved'
print(frame)

Ожидаемый вывод:

   id group_id    status  related_id
0   1        A   pending         NaN
1   2        A    active         NaN
2   3        A  resolved         1.0
3   4        B   pending         NaN
4   5        B    active         NaN
5   6        B  resolved         4.0
6   7        B  resolved         4.0

Какой подход выбрать

Оба подхода векторные и обходятся без groupby.apply. Выбирайте вариант с transform, если правило такое: «related_id относится к той же группе, а в группе есть хотя бы одна активная строка». Предпочтительнее merge, если правило таково: «решение принимается по статусу строки, на которую прямо указывает related_id в рамках той же группы».

Почему это важно

Опора на groupby.apply для межстрочных проверок часто вызывает ненужные замедления. Связка transform, map и merge оставляет вычисления на оптимизированном векторном пути pandas. Это лучше масштабируется и понятнее, когда условия завязаны на принадлежность к группе и поиск по другим строкам.

Выводы

Чётко формулируйте правило. Если условие — это существование на уровне группы (например, «в группе есть любой active») вместе со структурным ограничением («related_id в той же группе»), комбинируйте transform и map. Если нужно проверить конкретную строку, на которую указывает related_id, используйте merge — это прямой и эффективный путь. В обоих случаях вы избегаете построчного apply, сохраняете код лаконичным и получаете предсказуемую производительность.