PandasMiddleCoding

Как использовать query() для фильтрации DataFrame и когда он читабельнее булевой индексации?

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

DataFrame.query() — синтаксис

query() принимает строку с условием на языке, похожем на Python-выражения, и возвращает отфильтрованный DataFrame. Внутри строки можно использовать имена столбцов напрямую без df["col"].

import pandas as pd
import numpy as np

df = pd.DataFrame({
    "name": ["Alice", "Bob", "Carol", "Dave", "Eve"],
    "age": [25, 32, 28, 45, 22],
    "salary": [60000, 85000, 72000, 120000, 55000],
    "department": ["Eng", "Sales", "Eng", "Mgmt", "Sales"],
    "active": [True, True, False, True, False],
})

# Простое условие
result = df.query("age > 25 and salary >= 70000")

# Эквивалент булевой индексации:
result2 = df[(df["age"] > 25) & (df["salary"] >= 70000)]

# Оператор in
result3 = df.query("department in ['Eng', 'Mgmt']")

# not in
result4 = df.query("department not in ['Sales']")

# Внешняя переменная через @
min_salary = 70000
depts = ["Eng", "Sales"]
result5 = df.query("salary > @min_salary and department in @depts")

# Сравнение со строкой
result6 = df.query("name == 'Alice' or name == 'Bob'")

# Числовые выражения
result7 = df.query("salary / age > 2500")

# Работа с индексом
df_indexed = df.set_index("name")
result8 = df_indexed.query("index in ['Alice', 'Carol']")

# Метод str и регулярные выражения — через булевую индексацию
# (query не поддерживает .str напрямую):
result9 = df[df["name"].str.startswith("A")]

Когда query() читабельнее булевой индексации

  • Много условий: query("a > 1 and b < 5 and c == 'X'") vs. df[(df["a"]>1) & (df["b"]<5) & (df["c"]=="X")] — строка короче и понятнее.
  • Имена столбцов без повторения df: не нужно писать df["col"] в каждом условии.
  • Динамически построенные запросы: строку легко собрать из частей, что удобно при параметризованной фильтрации.
  • Читаемость в Jupyter-ноутбуках: цепочки df.query(...).groupby(...).agg(...) выглядят опрятно.

Производительность: numexpr

query() автоматически использует библиотеку numexpr, если она установлена и DataFrame достаточно большой (>10 000 строк). Numexpr компилирует выражение в байткод и вычисляет его поэлементно, используя все ядра CPU и меньше памяти (нет промежуточных массивов).

# pip install numexpr
import numexpr

# query автоматически переключится на numexpr при большом df
large_df = pd.DataFrame(np.random.randn(1_000_000, 5), columns=list("abcde"))

# Это быстрее обычной булевой индексации для больших данных:
result = large_df.query("a > 0.5 and b < -0.3")

# Принудительно отключить numexpr:
result2 = large_df.query("a > 0.5", engine="python")

Ограничения query()

  • Имена столбцов с пробелами или спецсимволами нужно оборачивать в обратные кавычки: df.query("`first name` == 'Alice'").
  • Не поддерживает методы строк (.str.contains), datetime-методы (.dt.year) — для этого нужна булевая индексация или eval().
  • Нельзя использовать pd.NA или np.nan напрямую в строке — проверяйте null через .isna() отдельно.

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

  • Имена столбцов с пробелами: без обратных кавычек запрос бросит SyntaxError или неочевидный UndefinedVariableError.
  • Конфликт имён столбцов с Python keywords: столбец not, and, in, True создаёт проблемы — переименуйте перед query.
  • Переменные окружения (@) работают только в локальной области видимости: если вызвать query внутри функции с переменной из внешней области — @var не найдёт её.
  • numexpr не поддерживает все операции: при использовании нестандартного синтаксиса Pandas тихо откатывается на Python engine без предупреждения.
  • query() не работает с MultiIndex столбцами: при MultiIndex нужна обычная булевая индексация.
  • Малые DataFrame: для DataFrame с <1000 строк numexpr не задействуется, а парсинг строки добавляет overhead — булевая индексация быстрее.
  • Изменение результата: query возвращает view или copy в зависимости от версии Pandas и CoW-настроек — не назначайте значения в результат query напрямую.

Common mistakes

  • Объяснять query filtering только синтаксисом без shape, dtype, состояния или режима выполнения.
  • Игнорировать leakage, воспроизводимость, пустые входы и скрытые копии данных.
  • Не проверять production-симптомы: latency, память, ретраи, дрейф качества и несовпадение версий.

What the interviewer is testing

  • Может ли связать query filtering с реальным контрактом входов и выходов.
  • Упоминает ли тесты, метрики, reproducibility и диагностику ошибок.
  • Видит ли различие между demo-кодом в ноутбуке и production-пайплайном.

Sources

Related topics