apache-spark - 将功能应用于Spark中的CSV的单个列

Question

使用Spark我正在读取一个csv，并想将一个函数应用于csv上的列。我有一些有效的代码，但它很hacky。正确的方法是什么？

我的密码

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.withColumn("message", udf_myFunction("_3")) #"_3" being the column name of the column you want to consider

这会将新列添加到包含 df结果的数据框 myFunction(line[3])中。