Agrupa y agrega dos columnas en un diccionario en Polars
Dado los siguientes datos, estoy buscando agrupar y combinar dos columnas en una, manteniendo un diccionario. Una columna proporciona las claves, mientras que los valores provienen de otra columna que se agrega en una lista primero.
import polars as pl
data = pl.DataFrame(
{
"names": ["foo", "ham", "spam", "cheese", "egg", "foo"],
"dates": ["1", "1", "2", "3", "3", "4"],
"groups": ["A", "A", "B", "B", "B", "C"],
}
)
print(data)
names dates groups
0 foo 1 A
1 ham 1 A
2 spam 2 B
3 cheese 3 B
4 egg 3 B
5 foo 4 C
Esto es lo que estoy intentando hacer:
Grupos Combinado
0 A {‘1’: [‘foo’, ‘ham’]}
1 B {‘2’: [‘spam’], ‘3’: [‘cheese’, ‘egg’]}
2 C {‘4’: [‘foo’]}
En pandas puedo hacer esto usando dos sentencias groupby, en pyspark utilizando un conjunto de operaciones alrededor de “mapfromentries”, pero a pesar de varios intentos no he encontrado una forma en polars. Hasta ahora utilizo agg_list(), convierto a pandas y uso una lambda. Si bien esto funciona, ciertamente no se siente correcto.
data = data.groupby(["groups", "dates"])["names"].agg_list()
data = (
data.to_pandas()
.groupby(["groups"])
.apply(lambda x: dict(zip(x["dates"], x["names_agg_list"])))
.reset_index(name="combined")
)
Alternativamente, inspirado por esta publicación, he intentado varias variaciones similares a la siguiente, incluida la conversión del dict a cadenas json, entre otras cosas.
data = data.groupby(["groups"]).agg(
pl.apply(exprs=["dates", "names_agg_list"], f=build_dict).alias("combined")
)
davy.ai
Para agrupar y combinar dos columnas en una, que contengan un diccionario, en Polars, puedes usar el método
agg_dict
junto con el métodoagg_list
para agrupar los datos porgroups
ydates
. Luego, puedes usar el métodoapply
con una función lambda para crear un diccionario a partir de las columnas agrupadas. Aquí tienes un ejemplo:Esto mostrará:
Aquí, primero agrupamos los datos por
groups
ydates
, y agregamos la columnanames
en una lista. Luego, usamos el métodoagg_dict
para crear un diccionario a partir dedates
ynames_agg_list
. Como el métodoagg_dict
devuelve un DataFrame con dos columnas (dates
ynames_agg_list
), usamos una función lambda conapply
para crear un diccionario a partir de estas columnas. Finalmente, eliminamos la columnadates
y ordenamos los datos porgroups
.