Упорядоченный Categorical в pandas без лишних преобразований

Как в pandas сразу получить упорядоченный категориальный столбец: пометить точные числовые ключи, остальные — как transient, без лишних операций и конверсий.

Помечать только отдельные числовые значения, а всё остальное относить к одной запасной метке — типичная задача при подготовке данных. Сложность начинается, когда итог должен быть не строками, а упорядоченным CategoricalDtype, да ещё и на таком объёме, где лишние преобразования превращаются в узкое место. Ниже — компактный приём, который обходит промежуточное сопоставление строк и сразу даёт упорядоченную категориальную колонку.

Проблема

Нужно присвоить читаемые метки нескольким точным значениям и собрать все прочие в единую категорию «transient», при этом сохранить результат как упорядоченный категориальный тип. Простые варианты через mapping и последующее приведение к категориальному работают, но выглядят громоздко и создают лишние накладные расходы на больших DataFrame.

import pandas as pd
MARKERS = [0, 9, 15, 25, 40]
LABELS = 'B BC/2 BC AB ABC'.split()
ordered_spec = pd.CategoricalDtype(categories=['transient', 'B', 'BC/2', 'BC', 'AB', 'ABC'], ordered=True)
frame = pd.DataFrame(data=[-3, 1.99, 0, 3, 9, 12, 15, 17, 24.9999999, 25, 25.000000001, 34, 40-1e-13, 40], columns=['value'])
frame['group'] = frame['value'].replace(dict(zip(MARKERS, LABELS)))
mask_other = ~frame['group'].isin(LABELS)
frame.loc[mask_other, 'group'] = 'transient'
frame['group'] = frame['group'].astype(ordered_spec)
print(frame)

Что происходит и почему это непросто

Нам нужно пометить только точные ключи, а все остальные числа оставить как «transient». Это несложно. Подвох в том, что финальный столбец должен иметь упорядоченный CategoricalDtype. Если сначала сопоставлять к строкам, а потом конвертировать в категориальный, вы делаете лишнюю работу и создаёте больше временных объектов, чем нужно — на больших объёмах это заметно. Вопрос в том, как получить упорядоченный категориальный сразу.

Решение

Постройте категориальный столбец поверх исходных числовых ключей. Приведите числовую Series к CategoricalDtype, где среди категорий присутствуют метка transient и точные ключи; затем переименуйте категории в нужные названия и заполните пропуски меткой transient. Так вы сразу получаете упорядоченный категориальный столбец.

default_tag = 'transient'
keyed_dtype = pd.CategoricalDtype(categories=[default_tag] + MARKERS, ordered=True)
frame['group'] = (
    frame['value']
        .astype(keyed_dtype)
        .cat.rename_categories([default_tag] + LABELS)
        .fillna(default_tag)
)

Нужен переиспользуемый помощник для произвольных входов? Оберните логику в небольшую функцию и пропустите через неё Series.

def build_ordered_cat(series_obj, key_values, name_values, default_tag='transient'):
    keyed_dtype = pd.CategoricalDtype(categories=[default_tag] + key_values, ordered=True)
    return (
        series_obj
            .astype(keyed_dtype)
            .cat.rename_categories([default_tag] + name_values)
            .fillna(default_tag)
    )
frame['group'] = build_ordered_cat(frame['value'], MARKERS, LABELS)

В итоге dtype — упорядоченный категориальный с ожидаемым порядком категорий:

frame['group'].dtype
# CategoricalDtype(categories=['transient', 'B', 'BC/2', 'BC', 'AB', 'ABC'], ordered=True, categories_dtype=object)

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

Этот подход избегает этапа со строками и последующей конвертации в категориальный. На больших DataFrame он по скорости сопоставим с map и заметно быстрее, чем связка map + fillna + astype(CategoricalDtype), потому что столбец изначально делается категориальным и не требует лишних преобразований. Плюс вы сохраняете предсказуемый, явно заданный порядок категорий на всём пути.

Итоги

Когда нужен упорядоченный категориальный столбец, который помечает только выбранные числовые ключи, задайте CategoricalDtype по этим ключам, выполните одно приведение, переименуйте категории в нужные подписи и заполните остальные значения единой запасной меткой, например «transient». Приём остаётся чистым, предсказуемым и эффективным на больших данных. Явно фиксируйте порядок категорий, делайте преобразование как можно раньше и избегайте лишних заходов через промежуточные строковые представления.