Цепочки методов

Анализ данных — это последовательность шагов: загрузить, добавить столбец, отфильтровать, сгруппировать, отсортировать. Код можно написать так, что эта последовательность видна с первого взгляда, — а можно так, что через неделю не разберётесь сами.

Промежуточные переменные

Три лучших сорта кофе по выручке. Первый способ — шаг за шагом, каждый результат в свою переменную:

import pandas as pd

orders = pd.read_csv("data/shop_orders.csv", parse_dates=["date"])
df = orders.copy()
df["revenue"] = df["price"] * df["quantity"]
df2 = df[df["category"] == "Кофе"]
df3 = df2.groupby("product")["revenue"].sum()
df4 = df3.sort_values(ascending=False)
df4.head(3)

product
Эфиопия 250 г          426600
Бразилия 1 кг          375390
Эспрессо-смесь 1 кг    372650
Name: revenue, dtype: int64

Работает. Но имена df2, df3, df4 ничего не значат, и за ними нужно следить: стоит перепутать df2 и df3 — и ошибка будет молчаливой. Переменные остаются в памяти, хотя больше не нужны. А если шаг нужно убрать, придётся переименовывать всё, что ниже.

Цепочка

Почти каждый метод pandas возвращает новую таблицу или Series — значит, следующий метод можно вызвать сразу у результата:

top_coffee = (
    orders
    .assign(revenue=orders["price"] * orders["quantity"])
    .query("category == 'Кофе'")
    .groupby("product")["revenue"]
    .sum()
    .sort_values(ascending=False)
    .head(3)
)
top_coffee

product
Эфиопия 250 г          426600
Бразилия 1 кг          375390
Эспрессо-смесь 1 кг    372650
Name: revenue, dtype: int64

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

Промежуточных переменных нет; лишний шаг удаляется или отключается одним символом # в начале строки.

top_cities = (
    orders
    # допишите шаги цепочки: каждый с новой строки, начиная с точки
)

top_cities = (
    orders
    .assign(revenue=orders["price"] * orders["quantity"])
    .groupby("city")["revenue"]
    .sum()
    .sort_values(ascending=False)
    .head(3)
)

lambda: таблица «на этом шаге»

В примере выше assign обращается к orders — исходной таблице. Это работает, пока assign стоит первым. Но если перед ним строки уже отфильтрованы или добавлен новый столбец, в orders этих изменений нет:

(
    orders
    .query("category == 'Чай'")
    .assign(revenue=orders["price"] * orders["quantity"])
    .assign(share=orders["revenue"] / orders["revenue"].sum())
    .head(3)
)

KeyError: 'revenue'

orders["revenue"] — ошибка: столбец revenue появился внутри цепочки, а orders остался прежним. У таблицы «на этом шаге» нет имени — к ней нельзя обратиться.

Для этого assign принимает не только готовые значения, но и функцию: pandas вызовет её и передаст текущую таблицу. Короткую функцию без имени пишут словом lambda:

(
    orders
    .query("category == 'Чай'")
    .assign(revenue=lambda d: d["price"] * d["quantity"])
    .assign(share=lambda d: d["revenue"] / d["revenue"].sum())
    [["product", "revenue", "share"]]
    .head(3)
)

Предпоследняя строка цепочки — [["product", "revenue", "share"]] — знакомый выбор столбцов. Он тоже шаг: квадратные скобки можно ставить сразу после предыдущего результата, без точки.

lambda d: d["price"] * d["quantity"] читается так: «функция, которая получает таблицу d и возвращает произведение её столбцов». Это то же самое, что

def make_revenue(d):
    return d["price"] * d["quantity"]

только короче и без имени. Букву после lambda выбираете вы; d или df — привычные варианты.

Правило: внутри цепочки обращайтесь к таблице через lambda — тогда шаги можно переставлять.

t = pd.DataFrame({"x": [1, 2, 3, 4]})
print(t.query("x > 2").assign(y=lambda d: d["x"] * 10)["y"].sum())

app_bulk = (
    orders
    # допишите шаги цепочки: каждый с новой строки, начиная с точки
)

app_bulk = (
    orders
    .query("channel == 'приложение' and quantity >= 3")
    .assign(revenue=lambda d: d["price"] * d["quantity"])
    .groupby("category")["revenue"]
    .sum()
    .sort_values(ascending=False)
)

pipe: свой шаг в цепочке

Часть шагов повторяется из анализа в анализ: добавить выручку, добавить месяц, взять первые N. Их удобно оформить функциями. Метод pipe («труба») вставляет функцию в цепочку: таблица передаётся ей первым аргументом, остальные аргументы пишут после названия функции.

def add_revenue(table):
    return table.assign(revenue=table["price"] * table["quantity"])


def top_rows(table, column, n):
    return table.nlargest(n, column)


(
    orders
    .pipe(add_revenue)
    .pipe(top_rows, "revenue", 3)
    [["date", "product", "revenue"]]
)

.pipe(top_rows, "revenue", 3) — это то же, что top_rows(таблица, "revenue", 3). Но запись с pipe сохраняет порядок чтения сверху вниз, а вложенные вызовы top_rows(add_revenue(orders), "revenue", 3) читаются изнутри наружу.

У функций-шагов одно правило: получить таблицу — вернуть новую таблицу, не меняя полученную.

def add_month(table):
    return ...


by_month = (
    orders
    # допишите шаги цепочки: каждый с новой строки, начиная с точки
)

def add_month(table):
    return table.assign(month=table["date"].dt.month)


by_month = (
    orders
    .pipe(add_revenue)
    .pipe(add_month)
    .groupby("month")["revenue"]
    .sum()
)

Когда цепочка не нужна

Цепочка хороша, пока читается. Если в ней пятнадцать шагов, разбейте её на две-три с осмысленными именами: priced, by_month. Если шаг упал с ошибкой — закомментируйте нижние строки и посмотрите на промежуточный результат. Цель — не «одна строка», а код, который легко прочитать.

Цикл по строкам

Тем, кто пришёл из обычного Python, хочется пройти по таблице циклом. Метод iterrows это позволяет: на каждом шаге даёт метку строки и саму строку как Series.

total = 0
for label, row in orders.head(3).iterrows():
    revenue = row["price"] * row["quantity"]
    print(label, row["product"], revenue)
    total = total + revenue
print(total)

0 Кофемолка ручная 6400
1 Травяной сбор 50 г 260
2 Шоколад горький 180
6840

Посчитаем циклом сумму к оплате со скидкой 10 % на позиции от трёх штук:

total = 0
for label, row in orders.iterrows():
    revenue = row["price"] * row["quantity"]
    if row["quantity"] >= 3:
        revenue = revenue * 0.9
    total = total + revenue
print(round(total))

3150009

Результат верный. Но так писать не стоит, и причин три:

• медленно: на каждой строке pandas собирает отдельный Series, а арифметику выполняет Python — для таблицы в миллион строк это десятки секунд вместо сотых долей секунды, разница в тысячи раз;
• длинно: шесть строк и переменная-накопитель вместо двух выражений;
• хрупко: результат цикла — одно число, а промежуточного столбца, который можно проверить глазами, нет.

Почти всё, что хочется сделать циклом, выражается операциями со столбцами: арифметика, np.where, map, groupby. Цикл по строкам оправдан редко — например, когда для каждой строки нужно отправить запрос в интернет.

import numpy as np

discounted = (
    orders
    # допишите шаги цепочки: каждый с новой строки, начиная с точки
)
total_fast = ...

import numpy as np

discounted = (
    orders
    .assign(revenue=lambda d: d["price"] * d["quantity"])
    .assign(to_pay=lambda d: np.where(d["quantity"] >= 3, d["revenue"] * 0.9, d["revenue"]))
)
total_fast = discounted["to_pay"].sum()

Итоги

• Методы возвращают новую таблицу — их можно вызывать цепочкой; всё выражение — в круглых скобках, шаг на строку.
• В assign внутри цепочки пишите lambda d: ... — d будет таблицей на этом шаге.
• pipe(функция, аргументы) вставляет свою функцию как шаг.
• Цикл iterrows работает, но медленно; сначала ищите операцию со столбцами.