Категориальный тип — pandas

Коротко

category хранит набор допустимых значений (categories) и для каждой строки — номер значения (codes).
Для столбцов вроде города, размера, статуса это в разы уменьшает память и ускоряет группировку.
ordered=True задаёт порядок: сортировка и сравнения идут по нему, а не по алфавиту.
Записать можно только значение из набора категорий; новое сначала добавляют через cat.add_categories.
Категории без строк видны в value_counts, но не в groupby: там с pandas 3.0 по умолчанию observed=True.

Синтаксис

s.astype("category")
pd.CategoricalDtype(["S", "M", "L"], ordered=True)
s.cat.categories, s.cat.codes
s.cat.add_categories([...]), s.cat.remove_unused_categories()
s.cat.rename_categories({...})

Параметры

ПараметрЧто делаетПо умолчанию

categoriesдопустимые значения и их порядокуникальные значения по алфавиту

orderedесть ли порядок: сравнения < > и сортировка по немуFalse

Примеры

Данные для примеровиспользуются во всех примерах ниже

sizes = pd.Series(["M", "S", "L", "M", "M", "S"], name="size")

Категории и коды

cat = sizes.astype("category")
print(cat.cat.categories)            # набор допустимых значений
cat.cat.codes.tolist()               # внутри хранятся номера

Вывод

Index(['L', 'M', 'S'], dtype='str')
[1, 2, 0, 1, 1, 2]

Упорядоченные категории

order = pd.CategoricalDtype(["S", "M", "L"], ordered=True)
ordered = sizes.astype(order)
print(ordered.sort_values().tolist())   # по заданному порядку, а не по алфавиту
ordered[ordered > "S"].tolist()          # сравнения по порядку

Вывод

['S', 'S', 'M', 'M', 'M', 'L']
['M', 'L', 'M', 'M']

Экономия памяти

big = pd.Series(["Москва", "Санкт-Петербург", "Казань"] * 100_000)
before = big.memory_usage(deep=True, index=False)
after = big.astype("category").memory_usage(deep=True, index=False)
print(f"{before / 2**20:.1f} МБ → {after / 2**20:.1f} МБ")
before // after                      # во сколько раз меньше

Вывод

7.3 МБ → 0.3 МБ
25

Пустые категории в value_counts и groupby

cat = sizes.astype(pd.CategoricalDtype(["S", "M", "L", "XL"]))
print(cat.value_counts())                          # value_counts показывает и пустую категорию
pd.DataFrame({"size": cat, "n": 1}).groupby("size").size()   # groupby — нет (observed=True)

Вывод

size
M     3
S     2
L     1
XL    0
Name: count, dtype: int64
size
S    2
M    3
L    1
dtype: int64

Подробнее — в подводных камнях статьи groupby.

Подводные камни

Новое значение не из набора

Запись нового значенияошибка: TypeError

cat = sizes.astype("category")
cat[0] = "XXL"                                     # значения нет среди категорий

Ошибка

TypeError: Cannot setitem on a Categorical with a new category (XXL), set the categories first

Сначала добавить категорию

cat = sizes.astype("category").cat.add_categories(["XXL"])
cat[0] = "XXL"
cat.tolist()

Вывод

['XXL', 'S', 'L', 'M', 'M', 'S']

Когда category не помогает

Если почти все значения разные (имена, адреса, id), категории хранят те же строки плюс номера — памяти уходит даже больше. Выигрыш заметен, когда уникальных значений намного меньше, чем строк.

Все значения разные

ids = pd.Series([f"id{i}" for i in range(10_000)])   # 10 000 разных строк
before = ids.memory_usage(deep=True, index=False)
after = ids.astype("category").memory_usage(deep=True, index=False)
after > before

Вывод

True