pyspark：如何在 spark dataframe 中對 N 條記錄進行分組

Question

我有一個 CSV 有 500 萬條記錄，結構如下：

+----------+------------+------------+
|  row_id  |    col1    |    col2    |
+----------+------------+------------+
|         1|   value    |    value   |
|         2|   value    |    value   |
|....                                |
|...                                 |
|   5000000|   value    |    value   |
+----------+------------+------------+

我需要將這個 CSV 轉換為 JSON，每個 json 文件都有 500 條記錄和一個特定的結構，如下所示：

{
    "entry": [
        {
            "row_id": "1",
            "col1": "value",
            "col2": "value"
        },
        {
            "row_id": "2",
            "col1": "value",
            "col2": "value"
        },
        ....
        ..
        {
            "row_id": "500",
            "col1": "value",
            "col2": "value"
        }
    ],
    "last_updated":"09-09-2021T01:03:04.44Z"
}

使用PySpark我能夠讀取 csv 並創建一個 dataframe。我不知道如何將 500 條記錄分組到結構"entry": [ <500 records> ],"last_updated":"09-09-2021T01:03:04.44Z"
我可以使用df.coalesce(1).write.option("maxRecordsPerFile",500)但這只會給我 500 條記錄的集合，沒有任何結構。 我想要"entry"列表中的那 500 條記錄和它后面的"last_updated" （我從datetime.now()中獲取）。

Answer 1

您可以嘗試以下操作：

---

注意。 我使用了以下導入。

from pyspark.sql import functions as F
from pyspark.sql import Window

---

1 . 我們需要一個可用於將您的數據拆分為 500 個記錄批次的列

（推薦）我們可以創建一個偽列來實現這一點row_number

df = df.withColumn("group_num",(F.row_number().over(Window.orderBy("row_id"))-1) % 500 )

否則，如果從1開始的row_id在 500 萬條記錄中持續增加，我們可以使用

df = df.withColumn("group_num",(F.col("row_id")-1) % 500 )

或者在這種奇怪的情況下， "last_updated":"09-09-2021T01:03:04.44Z"列對於每批 500 條記錄都是唯一的

df = df.withColumn("group_num",F.col("last_updated"))

---

2 . 我們將通過按group_num分組來轉換您的數據集

df = (
    df.groupBy("group_num")
      .agg(
          F.collect_list(
              F.expr("struct(row_id,col1,col2)")
          ).alias("entries")
      )
      .withColumn("last_updated",F.lit(datetime.now())))
      .drop("group_num")
)

注意。 如果您想包括所有列，您可以使用F.expr("struct(*)")而不是F.expr("struct(row_id,col1,col2)") 。

---

3 . 最后，您可以使用選項.option("maxRecordsPerFile",1)寫入輸出/目標，因為現在每行最多存儲 500 個條目

例如。

df.write.format("json").option("maxRecordsPerFile",1).save("<your intended path here>")

---

讓我知道這是否適合你