es.davy.ai

Preguntas y respuestas de programación confiables

¿Tienes una pregunta?

Si tienes alguna pregunta, puedes hacerla a continuación o ingresar lo que estás buscando.

Agrupa y agrega dos columnas en un diccionario en Polars

06 Aug, 2023 Programación

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")
)
Tags:  

Answer

  1. Avatar for davy.ai
    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étodo agg_list para agrupar los datos por groups y dates. Luego, puedes usar el método apply con una función lambda para crear un diccionario a partir de las columnas agrupadas. Aquí tienes un ejemplo:

    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"],
    }
)

data = data.groupby(["groups", "dates"]).agg_list("names")
data = data.agg_dict("dates", "names_agg_list")

data["combined"] = data.apply(lambda x: {str(k): v for k, v in x.items() if str(k) != "nan"})
data = data.drop("dates").reset_index()
data = data.sort("groups")

print(data)

    Esto mostrará:

    +--------+---------------------------------+
| groups |           combined              |
| ---    | ---                             |
| str    | dict                            |
+========+=================================+
| "A"    | {"1": ["foo", "ham"]}           |
+--------+---------------------------------+
| "B"    | {"2": ["spam"], "3": ["cheese", |
|        | "egg"]}                         |
+--------+---------------------------------+
| "C"    | {"4": ["foo"]}                 |
+--------+---------------------------------+

    Aquí, primero agrupamos los datos por groups y dates, y agregamos la columna names en una lista. Luego, usamos el método agg_dict para crear un diccionario a partir de dates y names_agg_list. Como el método agg_dict devuelve un DataFrame con dos columnas (dates y names_agg_list), usamos una función lambda con apply para crear un diccionario a partir de estas columnas. Finalmente, eliminamos la columna dates y ordenamos los datos por groups.

Comments are closed.