Split Dataframe to Dict of Dataframe using pyspark

Viewed 102

I have dataframe as below

id | Key   | Value |
-----------------------
0  | Key1  | 100   |
1  | Key1  | 101   |
2  | Key1  | 102   |
3  | Key1  | 103   |
4  | Key2  | 104   |
5  | Key2  | 105   |
6  | Key2  | 106   |
7  | Key3  | 107   |
8  | Key3  | 108   |
9  | Key3  | 109   |

I would like to split to dict of dataframe group by some column by using pyspark as below

{ "Key1" : id | Key   | Value |
          -----------------------
           0  | Key1  | 100   |
           1  | Key1  | 101   |
           2  | Key1  | 102   |
           3  | Key1  | 103   |,

  "Key2" : id | Key   | Value |
          -----------------------
           4  | Key2  | 104   |
           5  | Key2  | 105   |
           6  | Key2  | 106   |,

  "Key3" : id | Key   | Value |
          -----------------------
           7  | Key3  | 107   |
           8  | Key3  | 108   |
           9  | Key3  | 109   | }

I'm using spark 2.7.1 with pyspark.

I alredy try

out = dict()
for i in ["Key1", "Key2","Key3"]:
    out[i] = df.where(df.key == i)
return out

But i'm looking for other faster way

1 Answers

Even if the code scala the functions are same in Scala and PySpark

  • Group by Key
  • Create a map where key is Key and value is List of Struct this can be achieved using

    df.groupBy("Key") .agg( map(col("Key"),collect_list( struct(col("id"),col("Key"),col("Values")) )) .as("outputMap")) .show(false)

import org.apache.spark.sql.functions._

object GroupColumns {

  def main(args: Array[String]): Unit = {

    val spark = Constant.getSparkSess

    import spark.implicits._

    val df = List( (0,"Key1",100),(1,"Key1",101),(2,"Key1",102),(3,"Key1",103),
      (0,"Key2",100),(1,"Key2",101),(2,"Key2",102),(3,"Key2",103)).toDF("id","Key","Values")

//    df.show()

    df.groupBy("Key")
      .agg(
        map(col("Key"),collect_list( struct(col("id"),col("Key"),col("Values")) ))
      .as("outputMap"))
      .drop("Key")
      .show(false)
  }

}
Related