Применить функцию к одному столбцу csv в Spark

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].

Я использую Спарк версии 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 

Это добавит новый столбец в dataframe df, содержащей результат myFunction(line[3]).