在 Pyspark 中執行 foreachPartition 后獲取空數據幀

Question

我是 PySpark 的新手，我正在嘗試在我的數據幀中執行 foreachPartition 函數，然后我想使用相同的數據幀執行另一個函數。 問題是在使用 foreachPartition 函數后，我的數據框變空了，所以我無法用它做任何其他事情。 我的代碼如下所示：

def my_random_function(partition, parameters):
    #performs something with the dataframe
    #does not return anything

my_py_spark_dataframe.foreachPartition(
    lambda partition: my_random_function(partition, parameters))

有人能告訴我如何執行這個 foreachPartition 並使用相同的數據框來執行其他功能嗎？

我看到一些用戶談論使用 df.toPandas().copy() 復制數據框，但在我的例子中，這會導致一些執行問題，所以我想使用相同的數據框而不是創建一個新的數據框。

先感謝您！

Answer 1

不清楚您正在嘗試哪種操作； 但這是foreachPartition的示例用法：

示例數據是來自三大洲的國家/地區列表：

+---------+-------+
|Continent|Country|
+---------+-------+
|       NA|    USA|
|       NA| Canada|
|       NA| Mexico|
|       EU|England|
|       EU| France|
|       EU|Germany|
|     ASIA|  India|
|     ASIA|  China|
|     ASIA|  Japan|
+---------+-------+

以下代碼按“Continent”對數據進行分區，使用foreachPartition迭代每個分區，並將“Country”名稱寫入該特定分區的每個文件，即 continent。

df = spark.createDataFrame(data=[["NA", "USA"], ["NA", "Canada"], ["NA", "Mexico"], ["EU", "England"], ["EU", "France"], ["EU", "Germany"], ["ASIA", "India"], ["ASIA", "China"], ["ASIA", "Japan"]], schema=["Continent", "Country"])
df.withColumn("partition_id", F.spark_partition_id()).show()

df = df.repartition(F.col("Continent"))
df.withColumn("partition_id", F.spark_partition_id()).show()

def write_to_file(rows):
  for row in rows:
    with open(f"/content/sample_data/{row.Continent}.txt", "a+") as f:
      f.write(f"{row.Country}\n")

df.foreachPartition(write_to_file)

輸出：

三個文件：每個分區一個。

!ls -1 /content/sample_data/

ASIA.txt
EU.txt
NA.txt

每個文件都有該大陸（分區）的國家名稱：

!cat /content/sample_data/ASIA.txt
India
China
Japan

!cat /content/sample_data/EU.txt
England
France
Germany

!cat /content/sample_data/NA.txt
USA
Canada
Mexico