More efficient way to loop through PySpark DataFrame and create new columns
<p>I am converting some code written with Pandas to PySpark. The code has a lot of <code>for</code> loops to create a variable number of columns depending on user-specified inputs.</p>

<p>I'm using Spark 1.6.x, with the following sample code:</p>

<pre><code>from pyspark.sql import SQLContext
from pyspark.sql import functions as F
import pandas as pd
import numpy as np

# create a Pandas DataFrame, then convert to Spark DataFrame
test = sqlContext.createDataFrame(pd.DataFrame({'val1': np.arange(1,11)}))
</code></pre>

<p>Which leaves me with </p>

<pre><code>+----+
|val1|
+----+
|   1|
|   2|
|   3|
|   4|
|   5|
|   6|
|   7|
|   8|
|   9|
|  10|
+----+
</code></pre>

<p>I loop a lot in the code, for example the below:</p>

<pre><code>for i in np.arange(2,6).tolist():
    test = test.withColumn('val_' + str(i), F.lit(i ** 2) + test.val1)
</code></pre>

<p>Which results in: </p>

<pre><code>+----+-----+-----+-----+-----+
|val1|val_2|val_3|val_4|val_5|
+----+-----+-----+-----+-----+
|   1|    5|   10|   17|   26|
|   2|    6|   11|   18|   27|
|   3|    7|   12|   19|   28|
|   4|    8|   13|   20|   29|
|   5|    9|   14|   21|   30|
|   6|   10|   15|   22|   31|
|   7|   11|   16|   23|   32|
|   8|   12|   17|   24|   33|
|   9|   13|   18|   25|   34|
|  10|   14|   19|   26|   35|
+----+-----+-----+-----+-----+
</code></pre>

<p>**Question: ** How can I rewrite the above loop to be more efficient?</p>

<p>I've noticed that my code runs slower as Spark spends a lot of time on each group of loops (even on small datasets like 2GB of text input).</p>

<p>Thanks</p>

<p>One withColumn will work on entire rdd. So generally its not a good practise to use the method for every column you want to add. There is a way where you work with columns and their data inside a map function. Since one map function is doing the job here, the code to add new column and its data will be done in parallel. </p>

<p>a. you can gather new values based on the calculations</p>

<p>b. Add these new column values to main rdd as below</p>

<pre><code>val newColumns: Seq[Any] = Seq(newcol1,newcol2)
Row.fromSeq(row.toSeq.init ++ newColumns)
</code></pre>

<p>Here row, is the reference of row in map method</p>

<p>c. Create new schema as below</p>

<pre><code>val newColumnsStructType = StructType{Seq(new StructField("newcolName1",IntegerType),new StructField("newColName2", IntegerType))
</code></pre>

<p>d. Add to the old schema</p>

<pre><code>val newSchema = StructType(mainDataFrame.schema.init ++ newColumnsStructType)
</code></pre>

<p>e. Create new dataframe with new columns</p>

<pre><code>val newDataFrame = sqlContext.createDataFrame(newRDD, newSchema)
</code></pre>

python apache-spark pyspark