有效地從pyspark dataframe列創建新數據框

Question

我想知道在pyspark數據框中提取列並將其轉換為新數據框的最有效方法是什么？ 以下代碼在使用小型數據集時不會出現任何問題，但運行速度非常慢，甚至會導致內存不足錯誤。 我想知道如何提高這段代碼的效率嗎？

pdf_edges = sdf_grp.rdd.flatMap(lambda x: x).collect()  
edgelist = reduce(lambda a, b: a + b, pdf_edges, [])
sdf_edges = spark.createDataFrame(edgelist)

在pyspark dataframe sdf_grp中 ，“對”列包含以下信息

+-------------------------------------------------------------------+
|pairs                                                              |
+-------------------------------------------------------------------+
|[[39169813, 24907492], [39169813, 19650174]]                       |
|[[10876191, 139604770]]                                            |
|[[6481958, 22689674]]                                              |
|[[73450939, 114203936], [73450939, 21226555], [73450939, 24367554]]|
|[[66306616, 32911686], [66306616, 19319140], [66306616, 48712544]] |
+-------------------------------------------------------------------+

具有

root
|-- pairs: array (nullable = true)
|    |-- element: struct (containsNull = true)
|    |    |-- node1: integer (nullable = false)
|    |    |-- node2: integer (nullable = false)

我想將它們轉換為新的數據框sdf_edges如下所示

+---------+---------+
|    node1|    node2|
+---------+---------+
| 39169813| 24907492|
| 39169813| 19650174|
| 10876191|139604770|
|  6481958| 22689674|
| 73450939|114203936|
| 73450939| 21226555|
| 73450939| 24367554|
| 66306616| 32911686|
| 66306616| 19319140|
| 66306616| 48712544|
+---------+---------+

Answer 1

提取列的最有效方法是避免collect（） 。 當您調用collect（）時，所有數據都將傳輸到驅動程序並在那里進行處理。 要實現所需的更好方法是使用explode（）函數。 看下面的例子：

from pyspark.sql import types as T
import pyspark.sql.functions as F

schema = T.StructType([
  T.StructField("pairs", T.ArrayType(
      T.StructType([
          T.StructField("node1", T.IntegerType()),
          T.StructField("node2", T.IntegerType())
      ])
   )
   )
])


df = spark.createDataFrame(
[
([[39169813, 24907492], [39169813, 19650174]],),
([[10876191, 139604770]],        )                                    ,
([[6481958, 22689674]]      ,     )                                   ,
([[73450939, 114203936], [73450939, 21226555], [73450939, 24367554]],),
([[66306616, 32911686], [66306616, 19319140], [66306616, 48712544]],)
], schema)

df = df.select(F.explode('pairs').alias('exploded')).select('exploded.node1', 'exploded.node2')
df.show(truncate=False)

輸出：

+--------+---------+ 
|  node1 |   node2 | 
+--------+---------+ 
|39169813|24907492 | 
|39169813|19650174 | 
|10876191|139604770| 
|6481958 |22689674 | 
|73450939|114203936| 
|73450939|21226555 | 
|73450939|24367554 | 
|66306616|32911686 | 
|66306616|19319140 | 
|66306616|48712544 | 
+--------+---------+ 

Answer 2

好吧，我用下面的方法解決

sdf_edges = sdf_grp.select('pairs').rdd.flatMap(lambda x: x[0]).toDF()