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>There is a small overhead of repeatedly calling JVM method but otherwise for loop alone shouldn't be a problem. You can improve it slightly by using a single select:</p>

<pre><code>df = spark.range(1, 11).toDF("val1")

def make_col(i):
    return (F.pow(F.lit(i), 2) + F.col("val1")).alias("val_{0}".format(i))

spark.range(1, 11).toDF("val1").select("*", *(make_col(i) for i in range(2, 6)))
</code></pre>

<p>I would also avoid using NumPy types. Initializing NumPy objects is typically more expensive compared to plain Python objects and Spark SQL doesn't support NumPy types so there some additional conversions required.</p>

python apache-spark pyspark