将功能应用于火花中CSV的单列

Question

使用Spark我正在阅读CSV，并希望将功能应用于CSV上的列。我有一些可行的代码，但这很骇人听闻。这样做的正确方法是什么？

我的代码

SparkContext().addPyFile("myfile.py")
spark = SparkSession
    .builder
    .appName("myApp")
    .getOrCreate()
from myfile import myFunction
df = spark.read.csv(sys.argv[1], header=True,
    mode="DROPMALFORMED",)
a = df.rdd.map(lambda line: Row(id=line[0], user_id=line[1], message_id=line[2], message=myFunction(line[3]))).toDF()

我希望能够在列名上调用函数，而不是将每一行映射到line，然后在line[index]上调用该功能。

我正在使用Spark版本2.0.1

Answer 1

您可以简单地使用用户定义的功能（udf）与withColumn：

组合

from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf
udf_myFunction = udf(myFunction, IntegerType()) # if the function returns an int
df = df.withColumn("message", udf_myFunction("_3")) #"_3" being the column name of the column you want to consider

这将向包含myFunction(line[3])的结果的DataFrame df添加新列。