Creating a new dataframe from a pyspark dataframe column efficiently

Question

I wonder what is the most efficient way to extract a column in pyspark dataframe and turn them into a new dataframe? The following code runs without any problem with small datasets, but runs very slow and even causes out-of-memory error. I wonder how can I improve the efficiency of this code?

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

In pyspark dataframe sdf_grp , The column "pairs" contains information as below

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

with a schema of

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

I'd like to convert them into a new dataframe sdf_edges looks like below

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

Answer 1

The most efficient way to extract columns is avoiding collect() . When you call collect(), all the data is transfered to the driver and processed there. At better way to achieve what you want is using the explode() function. Have a look at the example below:

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)

Output:

+--------+---------+ 
|  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()