python - 使用 pyspark 将数据框拆分为数据框的字典

Question

我有如下数据框

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   |

我想通过使用 pyspark 将某些列拆分为数据帧组的字典，如下所示

{ "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   | }

我正在使用带有 pyspark 的 spark 2.7.1。

我已经尝试过了

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

但我正在寻找其他更快的方式

标签： pythonpandaspyspark

Answer 1

即使代码 scala 中的函数在 Scala 和 PySpark 中是相同的

• 通过...分组Key

• 创建一个mapwhere key is Keyand value is List of Struct 这可以使用来实现

  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)
  }

}