Выборка и фильтрация: урок по Python для анализа данных

Выборка строк и колонок встречается в каждой задаче на pandas, а на live-coding по ней сразу видно, писали вы реальный код или нет. Разберём, как выбирать, фильтровать и менять данные по условию без сюрпризов.

Колонки

df["price"]              # одна колонка, Series
df[["price", "qty"]]     # несколько колонок, DataFrame
df.filter(like="price")  # колонки, в названии которых есть "price"
df.select_dtypes("number")

Квадратные скобки у DataFrame выбирают колонки. Исключение: срез df[0:5] и булева маска df[mask] выбирают строки. Чтобы не путаться, для строк используйте loc и iloc.

loc и iloc

loc работает по меткам, iloc по позициям.

df.loc[10, "price"]           # строка с меткой 10, колонка price
df.loc[df["qty"] > 0, ["sku", "price"]]
df.iloc[0]                    # первая строка
df.iloc[:5, :2]               # первые 5 строк, первые 2 колонки

Главное отличие срезов: loc включает правую границу, iloc нет. На RangeIndex 0..n-1 метки и позиции совпадают, поэтому разница не видна, пока индекс не перемешается после фильтра или сортировки:

df = df.sort_values("price")
df.iloc[0]    # самая дешёвая строка
df.loc[0]     # строка с меткой 0, которая была первой ДО сортировки

Для одной ячейки есть быстрые df.at[label, "col"] и df.iat[i, j].

Булевы маски

Фильтр это Series из True и False той же длины:

mask = (df["city"] == "Казань") & (df["amount"] > 500)
df[mask]

Правила:

поэлементные операторы &, |, ~, а не and, or, not (они дают ValueError: The truth value of a Series is ambiguous);
каждое условие в скобках: & связывает сильнее, чем > и ==.

Полезные готовые маски:

df["status"].isin(["paid", "shipped"])
df["age"].between(18, 35)          # обе границы включены
df["email"].str.contains("@", na=False)
df["phone"].isna()

query

query читается как SQL-условие, переменные Python подставляются через @:

min_amount = 500
df.query("amount > @min_amount and city == 'Казань'")

Внутри query можно писать and/or, это отдельный парсер выражений. Удобно для длинных условий, но опечатки в именах колонок находятся только во время выполнения.

Изменение по условию

Только одним шагом через loc:

df.loc[df["amount"] > 1000, "discount"] = 0.1

Цепочное присваивание df[df["amount"] > 1000]["discount"] = 0.1 пишет в промежуточную копию. В pandas 3.0 с Copy-on-Write оно никогда не меняет df, в pandas 2.x работало непредсказуемо и давало SettingWithCopyWarning.

Для нескольких вариантов значения есть np.where и np.select:

import numpy as np

df["tier"] = np.select(
    [df["amount"] >= 10000, df["amount"] >= 1000],
    ["vip", "mid"],
    default="low",
)

np.select берёт первое истинное условие, поэтому порядок важен: сначала самое строгое.

Пропуски в условиях

Любое сравнение с NaN даёт False. Поэтому ~(df["x"] > 5) и df["x"] <= 5 не одно и то же: первое включит строки с пропуском, второе нет. Решайте судьбу пропусков явно через isna() или fillna() перед фильтром.

Топ и крайние значения

df.nlargest(10, "revenue")       # быстрее полной сортировки
df.nsmallest(3, "price")
df.loc[df["price"].idxmin()]     # строка с минимальной ценой

Частые ошибки

and вместо & и условия без скобок.
loc вместо iloc после сортировки или фильтра.
Цепочное присваивание вместо loc[маска, колонка].
Изменение ячеек в цикле for i in range(len(df)): медленно и ломается на индексе с дырами.
Забытые пропуски в отрицании условия.

Что спросят на собесе

Чем loc отличается от iloc, включает ли срез правую границу?
Почему and в фильтре падает и зачем скобки?
Как правильно изменить значения по условию, что такое цепочное присваивание?
Как задать значение из нескольких вариантов без apply?
Чем ~(x > 5) отличается от x <= 5?

Выборка и фильтрация в pandas: loc, iloc, маски