Agrupe por Id y obtenga múltiples registros para múltiples columnas en Scala.

Answer

1. 

   davy.ai

   May 31, 2023 at 1:14 am

   Para lograr el resultado deseado, necesitamos agrupar el marco de datos por id y luego aplicar ciertas condiciones basadas en el valor jerárquico mínimo presente en cada grupo. Esto puede implementarse utilizando las funciones groupBy, agg, when y otherwise de la siguiente manera:

   import org.apache.spark.sql.functions._
   
   val df_new = df.groupBy("id")
     .agg(
       when(min("hierarchy") === 1, 
         // select row with highest amount
         struct(lit(1) as "flag", max("amount") as "amount")
       )
       .when(min("hierarchy") === 2, 
         // select two rows with highest amount
         struct(lit(2) as "flag", max("amount") as "amount", 
           collect_list(struct("hierarchy", "amount")) as "hierarchy_list")
       )
       .otherwise(
         // select three rows with highest amount for hierarchy >= 4
         struct(lit(3) as "flag", max("amount") as "amount", 
           collect_list(struct("hierarchy", "amount")) as "hierarchy_list")
       ) 
       as "struct"
     )
     // explode hierarchy_list for flag=2 or 3, and sort by hierarchy and amount in descending order
     .selectExpr("id", "struct.amount", "explode_outer(struct.hierarchy_list) as (hierarchy, amount)")
     .orderBy(col("id"), col("struct.flag").desc, col("amount").desc)
     // group by id, flag and take first 3 values
     .groupBy(col("id"), col("struct.flag"))
     .agg(collect_list(struct(col("hierarchy"), col("amount"))).as("hierarchy_list"))
     .select(col("id"), explode_outer(col("hierarchy_list")).as(Seq("hierarchy", "amount")))
     .orderBy(col("id"), col("hierarchy").desc)
   

   Aquí, usamos la función when para aplicar diferentes condiciones basadas en el valor jerárquico mínimo en cada grupo. Para min("hierarchy") === 1, seleccionamos la fila con la cantidad más alta. Para min("hierarchy") === 2, seleccionamos las dos filas con la cantidad más alta y creamos una lista de valores de jerarquía y cantidad para jerarquía>=4. Para todos los demás casos, seleccionamos tres filas con la cantidad más alta para jerarquía>=4 y creamos una lista similar.

   A continuación, ordenamos el marco de datos resultante por id, flag (determinado por el valor jerárquico mínimo) y amount en orden descendente. Usamos explode_outer para desagrupar la lista de jerarquía para flag=2 o flag=3 y luego agrupamos el marco de datos por id, flag, y tomamos los primeros tres valores.

   Finalmente, seleccionamos las columnas id, hierarchy y amount y ordenamos el marco de datos por id y hierarchy en orden descendente.

   El marco de datos resultante df_new tendrá el resultado deseado como se muestra en la pregunta.

Comments are closed.