Как выбрать элементы в подгруппах groupby, или есть что то изящнее?

Question

Есть два фрейма - ордера с налогами и справочник налога на каждый товар.

Нужно отобрать по каждому коду товара те ордера, в которых значение "налога" выше стандартного, взятого из "справочника".

Вот код:

import pandas as pd


def marker(g, epsilon):
    """Mark."""
    g = g.query("fee + fee_preview < -@epsilon")
    return g


df_orders = pd.DataFrame({
    "order_id": list(range(1, 14)),
    "sku": 9 * ["a", ] + 4 * ["b", ],
    "fee_amount": [-1, -1, -1, -1, -1.1, -0.8, -1.1, -2, -2, -2.1, -2.1, -2.1, -4.3],
    "quanty": 7 * [1, ] + 2 * [2, ] + 3 * [1, ] + [2, ],
})

print(df_orders)

df_fee = pd.DataFrame({
    "fee": [1, 0.8, 2.1, 4.3],
    "sku": ["a", "d", "b", "g"],
})

print(df_fee)


df_orders["fee"] = df_orders["fee_amount"] / df_orders["quanty"]

df = df_orders.join(df_fee.set_index("sku"), on="sku", rsuffix="_preview")

dd = df.groupby("sku").apply(marker, 0.005)
print(dd)

Все отрабатывает, вот результат:

        order_id sku  fee_amount  quanty
0          1   a        -1.0       1
1          2   a        -1.0       1
2          3   a        -1.0       1
3          4   a        -1.0       1
4          5   a        -1.1       1
5          6   a        -0.8       1
6          7   a        -1.1       1
7          8   a        -2.0       2
8          9   a        -2.0       2
9         10   b        -2.1       1
10        11   b        -2.1       1
11        12   b        -2.1       1
12        13   b        -4.3       2
   fee sku
0  1.0   a
1  0.8   d
2  2.1   b
3  4.3   g
        order_id sku  fee_amount  quanty   fee  fee_preview
sku                                                        
a   4          5   a        -1.1       1 -1.10          1.0
    6          7   a        -1.1       1 -1.10          1.0
b   12        13   b        -4.3       2 -2.15          2.1

Но меня смущает операция .join(), раздувающая промежуточный фрейм.

Нет ли чего то "более SQL" подобного, или другого решения, боле экономно использующего память?

Answer 1

Львиная доля методов и функций в Pandas возвращают новый созданный результирующий DataFrame. Некоторые функции поддерживают параметр inplace=True - замену DataFrame "на месте" с экономией памяти. Но сами разработчики Pandas считают эту, казалось бы, полезную фичу злом и хотят от нее избавиться в будущих версиях.

Поэтому ваше решение вполне нормальное.

В том случае если вам нужно добавить только один новый столбец, объединяя по другому столбцу, можно воспользоваться методом Series.map():

df_orders["fee_preview"] = df_orders["sku"].map(df_fee.set_index("sku")["fee"])

но в любом случае лучше сравнить производительность обоих методов на немаленьких датасетах...

Answer 2

Я нашел решение. Вот как выглядит функция обработки группы:

def mark(g, epsilon, df_fee):
    """Select."""
    fee_preview = df_fee[df_fee.sku == g.sku.unique()[0]].fee.values[0]  # 1
    g = g.query("fee + @fee_preview < -@epsilon")                        # 2
    g["fee_preview"] = fee_preview                                       # 3
    return g

#1 - определяем sku группы
#2 - фильтрация нужных строк, почти без изменения как в исходном варианте
#3 - добавляем в группу столбец, в котором добавлен контрольный "налог" fee

и вызов теперь без join. Вместо

df = df_orders.join(df_fee.set_index("sku"), on="sku", rsuffix="_preview")
dd = df.groupby("sku").apply(marker, 0.005)

Теперь:

dd =  df_orders.groupby("sku").apply(select_0, 0.005, df_fee)

Лишний промежуточный фрейм df "уволен"...
Что интересно, столбец, добавляется на уровне каждой группы и потом объединяется pandas groupby автоматически. Рад, что так можно.

БЛОГ НА HUSL

Как выбрать элементы в подгруппах groupby, или есть что то изящнее?

Ответы (2 шт):