¿Cómo creo una nueva columna a partir de la salida de pandas groupby().sum()?

Intentando crear una nueva columna a partir de la groupby cálculo. En el código a continuación, obtengo los valores calculados correctos para cada fecha (ver el grupo a continuación), pero cuando intento crear una nueva columna (df['Data4']) con eso obtengo NaN. Así que estoy tratando de crear una nueva columna en el marco de datos con la suma de Data3 para todas las fechas y aplicar eso a cada fila de fecha. Por ejemplo, 2015-05-08 está en 2 filas (el total es 50+5 = 55) y en esta nueva columna me gustaría tener 55 en ambas filas.

import pandas as pd
import numpy as np
from pandas import DataFrame

df = pd.DataFrame({
    'Date' : ['2015-05-08', '2015-05-07', '2015-05-06', '2015-05-05', '2015-05-08', '2015-05-07', '2015-05-06', '2015-05-05'], 
    'Sym'  : ['aapl', 'aapl', 'aapl', 'aapl', 'aaww', 'aaww', 'aaww', 'aaww'], 
    'Data2': [11, 8, 10, 15, 110, 60, 100, 40],
    'Data3': [5, 8, 6, 1, 50, 100, 60, 120]
})

group = df['Data3'].groupby(df['Date']).sum()

df['Data4'] = group

quieres usar transform esto devolverá una Serie con el índice alineado con el df para que luego pueda agregarlo como una nueva columna:

In [74]:

df = pd.DataFrame({'Date': ['2015-05-08', '2015-05-07', '2015-05-06', '2015-05-05', '2015-05-08', '2015-05-07', '2015-05-06', '2015-05-05'], 'Sym': ['aapl', 'aapl', 'aapl', 'aapl', 'aaww', 'aaww', 'aaww', 'aaww'], 'Data2': [11, 8, 10, 15, 110, 60, 100, 40],'Data3': [5, 8, 6, 1, 50, 100, 60, 120]})
​
df['Data4'] = df['Data3'].groupby(df['Date']).transform('sum')
df
Out[74]:
   Data2  Data3        Date   Sym  Data4
0     11      5  2015-05-08  aapl     55
1      8      8  2015-05-07  aapl    108
2     10      6  2015-05-06  aapl     66
3     15      1  2015-05-05  aapl    121
4    110     50  2015-05-08  aaww     55
5     60    100  2015-05-07  aaww    108
6    100     60  2015-05-06  aaww     66
7     40    120  2015-05-05  aaww    121

¿Cómo creo una nueva columna con Groupby().Sum()?

Hay dos formas: una sencilla y otra un poco más interesante.

El favorito de todos: GroupBy.transform() con 'sum'

La respuesta de @Ed Chum se puede simplificar un poco. Llamada DataFrame.groupby en vez de Series.groupby. Esto da como resultado una sintaxis más simple.

# The setup.
df[['Date', 'Data3']]

         Date  Data3
0  2015-05-08      5
1  2015-05-07      8
2  2015-05-06      6
3  2015-05-05      1
4  2015-05-08     50
5  2015-05-07    100
6  2015-05-06     60
7  2015-05-05    120

df.groupby('Date')['Data3'].transform('sum')

0     55
1    108
2     66
3    121
4     55
5    108
6     66
7    121
Name: Data3, dtype: int64 

Es un poco más rápido,

df2 = pd.concat([df] * 12345)

%timeit df2['Data3'].groupby(df['Date']).transform('sum')
%timeit df2.groupby('Date')['Data3'].transform('sum')

10.4 ms ± 367 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
8.58 ms ± 559 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Poco convencional, pero digno de su consideración: GroupBy.sum() + Series.map()

Me topé con una idiosincrasia interesante en la API. Por lo que digo, puede reproducir esto en cualquier versión principal superior a 0.20 (probé esto en 0.23 y 0.24). Parece que puedes reducir constantemente unos pocos milisegundos del tiempo que tarda transform si en su lugar utiliza una función directa de GroupBy y transmitirlo usando map:

df.Date.map(df.groupby('Date')['Data3'].sum())

0     55
1    108
2     66
3    121
4     55
5    108
6     66
7    121
Name: Date, dtype: int64

Comparar con

df.groupby('Date')['Data3'].transform('sum')

0     55
1    108
2     66
3    121
4     55
5    108
6     66
7    121
Name: Data3, dtype: int64

Mis pruebas muestran que map es un poco más rápido si puede permitirse el lujo de usar el directo GroupBy función (como mean, min, max, first, etc.). Es más o menos más rápido para la mayoría de las situaciones generales hasta alrededor de ~200 mil registros. Después de eso, el rendimiento realmente depende de los datos.

(Izquierda: v0.23, Derecha: v0.24)

Buena alternativa para conocer, y mejor si tiene marcos más pequeños con un número menor de grupos. . . pero recomendaría transform como primera opción. Pensé que esto valía la pena compartirlo de todos modos.

Código de evaluación comparativa, para referencia:

import perfplot

perfplot.show(
    setup=lambda n: pd.DataFrame({'A': np.random.choice(n//10, n), 'B': np.ones(n)}),
    kernels=[
        lambda df: df.groupby('A')['B'].transform('sum'),
        lambda df:  df.A.map(df.groupby('A')['B'].sum()),
    ],
    labels=['GroupBy.transform', 'GroupBy.sum + map'],
    n_range=[2**k for k in range(5, 20)],
    xlabel="N",
    logy=True,
    logx=True
)

Sugiero en general usar el más poderoso applycon el que puede escribir sus consultas en expresiones individuales incluso para usos más complicados, como definir una nueva columna cuyos valores se definen como operaciones en grupos, y que también pueden tener diferentes valores dentro de el mismo grupo!

Esto es más general que el caso simple de definir una columna con el mismo valor para cada grupo (como sum en esta pregunta, que varía según el grupo es el mismo dentro del mismo grupo).

Caso simple (columna nueva con el mismo valor dentro de un grupo, diferente entre grupos):

# I'm assuming the name of your dataframe is something long, like
# `my_data_frame`, to show the power of being able to write your
# data processing in a single expression without multiple statements and
# multiple references to your long name, which is the normal style
# that the pandas API naturally makes you adopt, but which make the
# code often verbose, sparse, and a pain to generalize or refactor

my_data_frame = pd.DataFrame({
    'Date': ['2015-05-08', '2015-05-07', '2015-05-06', '2015-05-05', '2015-05-08', '2015-05-07', '2015-05-06', '2015-05-05'], 
    'Sym': ['aapl', 'aapl', 'aapl', 'aapl', 'aaww', 'aaww', 'aaww', 'aaww'], 
    'Data2': [11, 8, 10, 15, 110, 60, 100, 40],
    'Data3': [5, 8, 6, 1, 50, 100, 60, 120]})
​
(my_data_frame
    # create groups by 'Date'
    .groupby(['Date'])
    # for every small Group DataFrame `gdf` with the same 'Date', do:
    # assign a new column 'Data4' to it, with the value being
    # the sum of 'Data3' for the small dataframe `gdf`
    .apply(lambda gdf: gdf.assign(Data4=lambda gdf: gdf['Data3'].sum()))
    # after groupby operations, the variable(s) you grouped by on
    # are set as indices. In this case, 'Date' was set as an additional
    # level for the (multi)index. But it is still also present as a
    # column. Thus, we drop it from the index:
    .droplevel(0)
)

### OR

# We don't even need to define a variable for our dataframe.
# We can chain everything in one expression

(pd
    .DataFrame({
        'Date': ['2015-05-08', '2015-05-07', '2015-05-06', '2015-05-05', '2015-05-08', '2015-05-07', '2015-05-06', '2015-05-05'], 
        'Sym': ['aapl', 'aapl', 'aapl', 'aapl', 'aaww', 'aaww', 'aaww', 'aaww'], 
        'Data2': [11, 8, 10, 15, 110, 60, 100, 40],
        'Data3': [5, 8, 6, 1, 50, 100, 60, 120]})
    .groupby(['Date'])
    .apply(lambda gdf: gdf.assign(Data4=lambda gdf: gdf['Data3'].sum()))
    .droplevel(0)
)

Fuera:

(¿Por qué la expresión de python está entre paréntesis? Para que no tengamos que salpicar nuestro código con barras invertidas por todos lados, y podamos poner comentarios dentro de nuestro código de expresión para describir cada paso).

¿Qué tiene de poderoso esto? Es que está aprovechando todo el poder del “dividir-aplicar-combinar paradigma”. Le permite pensar en términos de “dividir su marco de datos en bloques” y “ejecutar operaciones arbitrarias en esos bloques” sin reducir/agregar, es decir, sin reducir el número de filas. (Y sin escribir bucles explícitos y detallados y recurrir a costosas uniones o concatenaciones para unir los resultados).

Consideremos un ejemplo más complejo. Uno en el que tiene varias series temporales de datos en su marco de datos. Tiene una columna que representa un tipo de producto, una columna que tiene marcas de tiempo y una columna que contiene la cantidad de artículos vendidos para ese producto en algún momento del año. Le gustaría agrupar por producto y obtener una nueva columna que contenga el total acumulado de los artículos que se venden para cada categoría. Queremos una columna que, dentro de cada “bloque” con el mismo producto, siga siendo una serie temporal y aumente monótonamente (solo dentro de un bloque).

¿Cómo podemos hacer esto? Con groupby + apply!

(pd
     .DataFrame({
        'Date': ['2021-03-11','2021-03-12','2021-03-13','2021-03-11','2021-03-12','2021-03-13'], 
        'Product': ['shirt','shirt','shirt','shoes','shoes','shoes'], 
        'ItemsSold': [300, 400, 234, 80, 10, 120],
        })
    .groupby(['Product'])
    .apply(lambda gdf: (gdf
        # sort by date within a group
        .sort_values('Date')
        # create new column
        .assign(CumulativeItemsSold=lambda df: df['ItemsSold'].cumsum())))
    .droplevel(0)
)

Fuera:

¿Otra ventaja de este método? ¡Funciona incluso si tenemos que agrupar por varios campos! Por ejemplo, si tuviéramos un 'Color' campo para nuestros productos, y queríamos la serie acumulativa agrupada por (Product, Color)podemos:

(pd
     .DataFrame({
        'Date': ['2021-03-11','2021-03-12','2021-03-13','2021-03-11','2021-03-12','2021-03-13',
                 '2021-03-11','2021-03-12','2021-03-13','2021-03-11','2021-03-12','2021-03-13'], 
        'Product': ['shirt','shirt','shirt','shoes','shoes','shoes',
                    'shirt','shirt','shirt','shoes','shoes','shoes'], 
        'Color': ['yellow','yellow','yellow','yellow','yellow','yellow',
                  'blue','blue','blue','blue','blue','blue'], # new!
        'ItemsSold': [300, 400, 234, 80, 10, 120,
                      123, 84, 923, 0, 220, 94],
        })
    .groupby(['Product', 'Color']) # We group by 2 fields now
    .apply(lambda gdf: (gdf
        .sort_values('Date')
        .assign(CumulativeItemsSold=lambda df: df['ItemsSold'].cumsum())))
    .droplevel([0,1]) # We drop 2 levels now

Fuera:

(Esta posibilidad de extenderse fácilmente a la agrupación en múltiples campos es la razón por la que me gusta poner los argumentos de groupby siempre en una lista, incluso si es un solo nombre, como ‘Producto’ en el ejemplo anterior).

Y puedes hacer todo esto sintéticamente en una sola expresión. (Claro, si las lambdas de python fueran un poco más agradables a la vista, se verían aún mejor).

¿Por qué repasé un caso general? Porque esta es una de las primeras preguntas SO que aparece al buscar en Google cosas como “pandas new column groupby”.

Pensamientos adicionales sobre la API para este tipo de operación

Agregar columnas basadas en cálculos arbitrarios realizados en grupos es muy parecido al buen idioma de definiendo una nueva columna usando agregaciones sobre Windows en SparkSQL.

Por ejemplo, puedes pensar en esto (es código Scala, pero el equivalente en PySpark se ve prácticamente igual):

val byDepName = Window.partitionBy('depName)
empsalary.withColumn("avg", avg('salary) over byDepName)

como algo así como (usando pandas en la forma que hemos visto anteriormente):

empsalary = pd.DataFrame(...some dataframe...)
(empsalary
    # our `Window.partitionBy('depName)`
    .groupby(['depName'])
    # our 'withColumn("avg", avg('salary) over byDepName)
    .apply(lambda gdf: gdf.assign(avg=lambda df: df['salary'].mean()))
    .droplevel(0)
)

(Observe cuánto más sintético y mejor es el ejemplo de Spark. El equivalente de pandas parece un poco torpe. La API de pandas no facilita la escritura de este tipo de operaciones “fluidas”).

Este modismo a su vez proviene de Funciones de ventana de SQLque la documentación de PostgreSQL da una muy buena definición de: (énfasis mío)

Una función de ventana realiza un cálculo a través de un conjunto de filas de la tabla que de alguna manera están relacionadas con la fila actual. Esto es comparable al tipo de cálculo que se puede hacer con una función agregada. Pero a diferencia de las funciones agregadas normales, el uso de una función de ventana no hace que las filas se agrupen en una sola fila de salida: las filas conservan sus identidades separadas. Detrás de escena, la función de ventana puede acceder a más que solo la fila actual del resultado de la consulta.

Y da un hermoso ejemplo de SQL de una sola línea: (clasificación dentro de los grupos)

SELECT depname, empno, salary, rank() OVER (PARTITION BY depname ORDER BY salary DESC) FROM empsalary;

Lo último: también te puede interesar pandas’ pipeque es similar a apply pero funciona un poco diferente y le da a las operaciones internas un mayor alcance para trabajar. Ver aquí para más

df = pd.DataFrame({
'Date' : ['2015-05-08', '2015-05-07', '2015-05-06', '2015-05-05', '2015-05-08', '2015-05-07', '2015-05-06', '2015-05-05'], 
'Sym'  : ['aapl', 'aapl', 'aapl', 'aapl', 'aaww', 'aaww', 'aaww', 'aaww'], 
'Data2': [11, 8, 10, 15, 110, 60, 100, 40],
'Data3': [5, 8, 6, 1, 50, 100, 60, 120]
})
print(pd.pivot_table(data=df,index='Date',columns="Sym",     aggfunc={'Data2':'sum','Data3':'sum'}))

producción

Data2      Data3     
Sym         aapl aaww  aapl aaww
Date                            
2015-05-05    15   40     1  120
2015-05-06    10  100     6   60
2015-05-07     8   60     8  100
2015-05-08    11  110     5   50