循环通过 PySpark DataFrame 并创建新列的更有效方法

Question

我正在将一些用 Pandas 编写的代码转换为 PySpark。该代码有很多 for 循环来根据用户指定的输入创建可变数量的列。

我使用的是 Spark 1.6.x，示例代码如下：

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)}))

这让我留下了

+----+
|val1|
+----+
|   1|
|   2|
|   3|
|   4|
|   5|
|   6|
|   7|
|   8|
|   9|
|  10|
+----+

我在代码中循环了很多，例如下面：

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

结果：

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

**问题：**如何重写上述循环以提高效率？

我注意到我的代码运行速度较慢，因为 Spark 在每组循环上花费了大量时间（即使是在像 2GB 文本输入这样的小型数据集上）。

谢谢

Answer 1

重复调用 JVM 方法的开销很小，但单独的 for 循环应该不是问题。您可以通过使用单个选择来稍微改进它：

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

我也会避免使用 NumPy 类型。与普通 Python 对象相比，初始化 NumPy 对象的成本通常更高，而且 Spark SQL 不支持 NumPy 类型，因此需要进行一些额外的转换。

Answer 2

一个 withColumn 将适用于整个 rdd。因此，通常对要添加的每一列都使用该方法不是一个好习惯。有一种方法可以在 map 函数中处理列及其数据。由于一个 map 函数在这里完成这项工作，因此添加新列及其数据的代码将并行完成。

一个。您可以根据计算收集新值

b.如下将这些新列值添加到主 rdd

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

这里的row，是map方法中row的引用

c。如下创建新架构

val newColumnsStructType = StructType{Seq(new StructField("newcolName1",IntegerType),new StructField("newColName2", IntegerType))

d。添加到旧架构

val newSchema = StructType(mainDataFrame.schema.init ++ newColumnsStructType)

e。使用新列创建新数据框

val newDataFrame = sqlContext.createDataFrame(newRDD, newSchema)