Как фильтровать DataFrame по списку значений: 3 метода pandas

Для кого эта статья:

• Аналитики данных и разработчики, работающие с pandas
• Студенты и учащиеся курсов по обработке данных

• Профессионалы, стремящиеся улучшить свои навыки в программировании и оптимизации кода

  Работа с DataFrame в pandas — это как археологические раскопки: нужные данные скрыты под слоями ненужной информации. Когда требуется извлечь строки, соответствующие списку конкретных значений, многие разработчики теряются в поисках элегантного решения. 🔍 Хватит писать громоздкие конструкции с множеством условий! Существуют три мощных метода, которые превратят ваш код из неуклюжего монстра в изящную балерину. Давайте разберемся, как мастерски фильтровать DataFrame по списку значений и почему это критически важно для производительности ваших проектов.

Хотите превратиться из любителя pandas в настоящего профессионала обработки данных? На курсе Профессия аналитик данных от Skypro вы не только освоите продвинутые техники фильтрации DataFrame, но и научитесь оптимизировать код для работы с гигабайтами данных. Студенты курса получают уникальные знания о внутренних механизмах pandas, которые не найти в документации, и выходят на новый уровень профессионализма за 9 месяцев.

3 эффективных метода фильтрации DataFrame по списку

Pandas — настоящая швейцарская ножовка в мире анализа данных, предлагающая разные инструменты для решения одной задачи. Когда речь идет о фильтрации DataFrame по списку значений (аналог SQL-операции IN), у нас есть три основных подхода: использование метода .isin(), индексация через .loc[] и применение метода .query().

Прежде чем углубляться в детали, давайте создадим тестовый DataFrame, с которым будем работать:

import pandas as pd
import numpy as np

# Создаём тестовый DataFrame
df = pd.DataFrame({
'id': range(1, 1001),
'category': np.random.choice(['A', 'B', 'C', 'D', 'E'], 1000),
'value': np.random.randint(1, 100, 1000),
'region': np.random.choice(['North', 'South', 'East', 'West'], 1000)
})

Допустим, нам нужно выбрать все строки, где значение столбца 'category' находится в списке ['A', 'C', 'E']. Как это сделать наиболее эффективно? 🤔

Метод .isin() – простая фильтрация строк по множеству значений

Метод .isin() — самый интуитивно понятный и часто используемый подход. Он проверяет, содержатся ли значения столбца в заданном списке, и возвращает булевую маску:

# Создаем список категорий для фильтрации
categories_to_select = ['A', 'C', 'E']

# Применяем фильтр с использованием .isin()
filtered_df = df[df['category'].isin(categories_to_select)]

# Выводим количество полученных строк
print(f"Получено {len(filtered_df)} строк из {len(df)}")

Метод .isin() особенно удобен, когда нужно быстро отфильтровать данные без сложной логики. Он работает со всеми типами данных и не требует дополнительных импортов или настроек.

Алексей, ведущий аналитик данных

На одном проекте мне пришлось анализировать логи посещений крупного интернет-магазина с миллионами записей. Клиент попросил отследить поведение пользователей, приходящих из определенных рекламных кампаний. У меня был список из 50 ID кампаний, и я сначала написал громоздкое условие с операторами OR:

filtered_logs = logs[(logs['campaign_id'] == ids[0]) | 
(logs['campaign_id'] == ids[1]) | 
... и так далее]

Код был ужасен, медленно работал и занимал 50 строк. Коллега увидел это безобразие и показал метод .isin():

filtered_logs = logs[logs['campaign_id'].isin(campaign_ids)]

Одна строка решила проблему, а скорость выполнения выросла в 8 раз! С тех пор .isin() — мой верный союзник в ежедневных задачах.

Преимущества метода .isin():

• Простота и читаемость кода
• Работает со всеми типами данных
• Хорошо оптимизирован для больших DataFrame
• Позволяет комбинировать условия с другими фильтрами

Недостатки:

• Для очень сложных условий может быть не так гибок, как .query()
• При работе с большими списками значений (тысячи элементов) может быть медленнее некоторых альтернативных подходов

Использование .loc[] для точной выборки данных из DataFrame

Метод .loc[] — мощный инструмент для индексации в pandas, который можно комбинировать с .isin() для создания более сложных фильтров. Он особенно полезен, когда нужно не только выбрать строки, но и указать, какие именно столбцы должны быть в результате:

# Используем .loc[] с .isin() для выбора строк и определенных столбцов
result = df.loc[df['category'].isin(categories_to_select), ['id', 'category', 'value']]

# Проверяем результат
print(result.head())

Главное преимущество .loc[] — возможность одновременно фильтровать строки и выбирать конкретные столбцы, что делает код более компактным и читаемым.

Преимущества метода .loc[]:

• Повышенная производительность при одновременной фильтрации строк и выборе столбцов
• Более ясное указание намерения в коде
• Позволяет использовать индексы DataFrame для более быстрого доступа
• Предотвращает создание промежуточных копий данных

Недостатки:

• Для простых фильтраций может выглядеть избыточным
• Менее интуитивен для новичков в pandas

Метод .query() – SQL-подобный синтаксис для фильтрации DataFrame

Для тех, кто привык к SQL, метод .query() предлагает интуитивно понятный синтаксис, позволяющий писать условия фильтрации в виде строк. Это особенно удобно для сложных условий:

# Преобразуем список в строку для использования в query
categories_string = "', '".join(categories_to_select)
categories_string = f"['{categories_string}']"

# Применяем метод .query()
query_result = df.query(f"category in {categories_string}")

# Проверяем результат
print(query_result.head())

Метод .query() особенно полезен, когда условия фильтрации формируются динамически или когда нужно применить сложную логику. Он использует движок выражений numexpr (если установлен), что может значительно ускорить выполнение на больших наборах данных.

Марина, руководитель отдела аналитики

В нашей компании аналитики тратили много времени на создание отчетов по продажам, которые требовали фильтрации товаров по сложным условиям. Код был заполнен множеством вложенных условий с использованием & и |, что делало его практически нечитаемым:

filtered = sales[
((sales['category'].isin(['Electronics', 'Appliances'])) & 
(sales['price'] > 1000) & 
(sales['date'] >= '2022-01-01')) | 
((sales['category'].isin(['Books', 'Toys'])) & 
(sales['discount'] > 0.2))
]

Когда я обнаружила метод .query(), мы переписали код:

filtered = sales.query(
"(category in ['Electronics', 'Appliances'] and price > 1000 and date >= '2022-01-01') or "
"(category in ['Books', 'Toys'] and discount > 0.2)"
)

Это не только сделало код более читаемым и поддерживаемым, но и позволило нам создавать шаблоны запросов, которые легко модифицировать. Время на разработку отчетов сократилось на 40%, а количество ошибок уменьшилось драматически.

Преимущества метода .query():

• SQL-подобный синтаксис, понятный даже неопытным пользователям pandas
• Возможность динамического формирования условий
• Потенциальное ускорение при использовании numexpr
• Лаконичная запись сложных условий

Недостатки:

• Для полного раскрытия потенциала требует установки дополнительного пакета numexpr
• Может быть медленнее на небольших наборах данных
• Необходимость экранирования кавычек в строковых условиях

Сравнение производительности методов фильтрации на больших данных

Теоретические рассуждения хороши, но давайте проведем эксперимент и сравним производительность всех трех методов на разных размерах DataFrame и разных размерах списков для фильтрации. 🚀

Создадим функцию для тестирования производительности:

import time
import pandas as pd
import numpy as np

def benchmark_filtering_methods(df_size, list_size):
# Создаем большой DataFrame
df = pd.DataFrame({
'id': range(1, df_size + 1),
'category': np.random.choice(list(map(str, range(100))), df_size),
'value': np.random.randint(1, 100, df_size)
})

# Создаем список для фильтрации
filter_list = [str(i) for i in range(list_size)]

# Тестируем .isin()
start = time.time()
result1 = df[df['category'].isin(filter_list)]
isin_time = time.time() – start

# Тестируем .loc[]
start = time.time()
result2 = df.loc[df['category'].isin(filter_list)]
loc_time = time.time() – start

# Тестируем .query()
filter_string = "', '".join(filter_list)
filter_string = f"['{filter_string}']"
start = time.time()
result3 = df.query(f"category in {filter_string}")
query_time = time.time() – start

return {
'isin_time': isin_time,
'loc_time': loc_time,
'query_time': query_time
}

# Запускаем тесты для разных размеров данных
results = []
for df_size in [10_000, 100_000, 1_000_000]:
for list_size in [5, 20, 50]:
result = benchmark_filtering_methods(df_size, list_size)
results.append({
'df_size': df_size,
'list_size': list_size,
'isin_time': result['isin_time'],
'loc_time': result['loc_time'],
'query_time': result['query_time']
})

Результаты тестирования можно представить в виде таблицы:

Результаты тестирования показывают интересную закономерность:

• На небольших и средних DataFrame (до 100,000 строк) методы .isin() и .loc[] показывают почти идентичную производительность и обгоняют .query()
• На больших DataFrame (1,000,000 строк и более) метод .query() становится быстрее, особенно при больших списках фильтрации
• Разница в производительности между .isin() и .loc[] минимальна во всех сценариях
• Размер списка для фильтрации влияет на производительность, но не меняет принципиально соотношение между методами

Исходя из полученных результатов, можно сформулировать рекомендации по выбору метода фильтрации:

• Для небольших и средних DataFrame (до 100,000 строк) — используйте .isin() или .loc[] как наиболее простые и достаточно быстрые методы
• Для больших DataFrame (более 1,000,000 строк) — рассмотрите .query(), особенно если у вас установлен numexpr
• Если читаемость кода важнее производительности — .query() даст более понятный код для сложных условий
• Если требуется одновременная фильтрация строк и выбор столбцов — .loc[] будет наиболее элегантным решением

Помните, что для максимальной производительности метода .query() рекомендуется установить библиотеку numexpr:

pip install numexpr

На практике часто выбор метода фильтрации определяется не только производительностью, но и стилем кодирования команды, читаемостью кода и конкретными требованиями проекта. 💼

Выбор подходящего метода фильтрации DataFrame по списку значений — это баланс между производительностью и удобством использования. Метод .isin() отлично подходит для повседневных задач благодаря своей простоте. Метод .loc[] предлагает более точный контроль над выборкой данных. А .query() становится незаменимым при работе с гигантскими массивами информации. Осваивая все три метода и понимая их сильные стороны, вы значительно расширяете свой арсенал инструментов для эффективной обработки данных. Главное правило успешного аналитика — выбирать подходящий инструмент для конкретной задачи, а не пытаться применить один метод ко всем сценариям. 📊