Используем query и eval в pandas

Если вы работаете в сфере data science на Python, то вы точно используете библиотеку pandas. В ней есть множество полезных методов для работы с табличными структурами, знакомство с которыми может сделать ваш код более эффективным и читабельным.

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

import pandas as pd
import numpy as np

data = []
for _ in range(100):
    data.append({
        "gender": "Male",
        "height": np.random.normal(178, 10),
        "age": np.random.uniform(20, 70)
    })
for _ in range(100):
    data.append({
        "gender": "Female",
        "height": np.random.normal(166, 8),
        "age": np.random.uniform(20, 70)
    })
df = (pd.DataFrame(data)
    # Используем sample для перемешивания данных
    .sample(frac=1.0, replace=False)
    .reset_index(drop=True)
 )

Ниже представлен полученный набор:

Теперь давайте выполним пару операции: возьмём всех женщин в возрасте от 20 до 30 лет, а затем найдём самых высоких и самых низких в этой группе:

(
    df[(df["gender"] == "Female") & (df["age"] >= 20) & (df["age"] <= 30)]["height"]
    .pipe(lambda x: [x.max(), x.min()])
)
# [174.23275599549675, 152.7714804940481]

Возможно, вы согласитесь, что это выражение немного перегружено. Есть ли какой-нибудь другой способ сделать то же самое?

Используем query() и eval()

query и eval — это два метода, которые анализируют строковые выражения для выполнения операций более естественным («питоническим»‎) способом. Метод query() предназначен для фильтрации датафрейма, в то время как eval() используется для вычислений и присваиваний. Пока что мы не будем вдаваться в подробности, а просто посмотрим, как приведённое выше выражение может быть переведено для использования query() и eval():

(
    df.query("gender == 'Female' and 20 <= age <= 30")
    .eval("height.max(), height.min()")
)
# [174.23275599549675, 152.7714804940481]

Уже лучше, не правда ли? В первом выражении было множество квадратных скобок, а это выражение гораздо более чистое и простое. Хотя методы query() и eval() применимы далеко не всегда, в некоторых ситуациях они могут быть очень удобны.

Операции в цепочку

Одна из особенностей pandas — это операции в цепочку. Поскольку большинство из них возвращает не исходный датафрейм, а его копию (если вы не используете inplace=True), вы можете выполнять набор операций последовательно в одном выражении. И часто нам нужно получить ссылку на вновь созданный объект, чтобы выполнить фильтрацию, присваивание и вычисления. В этом случае используется метод pipe(). Например:

## НЕПРАВИЛЬНЫЙ СПОСОБ
a = df[df["gender"] == "Male"].reset_index(drop=True).assign(age=df.age-10)
# > Это неверно, потому что в окончательном присваивании df.age 
# относится к исходному df, а не к df только с мужчинами и reset_index!

# Как это следует делать
b = (df[df["gender"] == "Male"].reset_index(drop=True) 
.pipe(lambda x: x.assign(age=x.age-10)))

# Это не то же самое
assert not a.equals(b)

# Давайте добавим больше фильтров
b = (df[df["gender"] == "Male"].reset_index(drop=True)
    .pipe(lambda x: x.assign(age=x.age-10))
    .pipe(lambda x: x[x["age"] > 30]))

В конечном выражении мы дважды используем метод pipe(), из-за чего код становится немного громоздким и может стать ещё хуже при использовании более длинных конструкций. Преимущество методов query() и eval() заключается в том, что они будут знать, что вы ссылаетесь на столбцы внутри текущего фрейма данных, и, следовательно, не будет никакой необходимости использовать pipe():

c = (df.query("gender == 'Male'")
    .reset_index(drop=True)
    .eval("age=age-10")
    .query("age > 30"))
assert b.equals(c)

Теперь, когда вы знаете некоторые преимущества использования query() и eval(), давайте разберём их специфику.

query

Метод query() используется для фильтрации датаферймов, то есть для выбора подмножества строк. Вычисленное выражение анализируется, а затем помещается внутрь .loc[<выражение>], а в случае многомерного ключа используется df[...].

Операторы

Операции, которые вы можете выполнять, включают в себя стандартные операции, используемые при фильтрации: &, |, <, >, >=, <=, ==, ~. Особыми из них являются:

• and, or работает как &, |, в чём мы убедились в предыдущих примерах.
• in работает, как .isin(...). Например, df.query("gender in ['Male', 'Female']) вернёт все экземпляры, где пол мужской или женский. Если напишем not in, то получим объекты, которых нет в списке. Также вместо целого списка можно использовать название столбца, что эквивалентно .isin(df[<столбец>]).
• == <list> работает так же, как in, а != <список> — как not in, когда правая часть является списком.
• not работает, как ~.

Обращение к элементам

Для обращения к столбцам просто используем их названия без кавычек. Если в столбце есть пробел или другие специальные символы, то используем обратные кавычки: `моя переменная`. А вот переменные, которые являются зарезервированными ключевыми словами Python, этим методом не поддерживаются. Для использования переменных вне датафрейма применяется знак @. Например, чтобы использовать локальную переменную a = 10, мы можем выполнить df.query("age > @a"). Можно также использовать математические функции, например:

from numpy import power
import numpy as np

# Работает
df.query("@power(age, 2) > 1000")

# Не работает
df.query("@np.power(age, 2) > 1000")

eval

Метод eval используется для выполнения присваиваний или произвольных вычислений. Применяется тот же тип операторов и ссылок, что и в query. Некоторые из операций, которые невозможно выполнить — это создание объектов, таких как словари и множества, и выполнение list/dict/set comprehension. Примеры:

from numpy import power

# Присваивание переменной значения
df.eval("age_2 = @power(age, 2)")

# для присваивания значений нескольким переменным мы должны делать это
# с новых строк.
# Обратите внимание, что в третьем присваивании используется
# sqrt без @, потому что существуют некоторые
# стандартные операции, которые можно использовать без @
df.eval("""
    age_2 = @power(age, 2)
    age_sqrt = age**(1/2)
    age_sqrt = sqrt(age)
""") 
# Возвращает df с 3 новыми столбцами: age, age_2 and age_sqrt
# Как всегда, добавляя inplace=True, мы изменим исходный датафрейм

df.eval("age+3") # Вернёт столбец age + 3
df.eval("[3, 2]") # Вернёт [3, 2]
df.eval("height.min()") # Вернёт минимальный рост

Чтобы лучше всё понять, просто попрактикуйтесь самостоятельно! Вы поймёте, что в некоторых случаях лучше придерживаться обычных операторов pandas, зато в других query() и eval() работают идеально.

Заключение

Методы query() и eval() помогут вам писать простой и лаконичный код, а с практикой ваша эффективность только увеличится. Хорошего дня :)