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

时间:2016-10-14 19:28:37

标签: python apache-spark pyspark

我正在将一些用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文本输入等小型数据集上)。

由于

2 个答案:

答案 0 :(得分:3)

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

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类型,因此需要一些额外的转换。

答案 1 :(得分:-1)

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

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

湾将这些新列值添加到主rdd,如下所示

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

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

℃。创建新架构,如下所示

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

d。添加到旧架构

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

即使用新列创建新数据框

val newDataFrame = sqlContext.createDataFrame(newRDD, newSchema)