PySpark中DataFrame的逐行操作

Question

如果有一个数据帧，并希望根据行的值对函数中的数据进行一些操作。

my_udf(row):
    threshold = 10
        if row.val_x > threshold
        row.val_x = another_function(row.val_x)
        row.val_y = another_function(row.val_y)
        return row
    else:
        return row

有谁知道如何将我的udf应用于数据帧？

Answer 1

如果您可以使用pyspark函数，最好不要使用UDF，如果您不能将＜code＞另一个_function

from pyspark.sql.types import *
import pyspark.sql.functions as psf

def another_function(val):
    ...

another_function_udf = psf.udf(another_function, [outputType()])

其中< code>outputType()是与< code > other _ function (< code > integer type()，< code>StringType()的输出相对应的pyspark类型...)

def apply_another_function(val):
    return psf.when(df.val_x > threshold, another_function_udf(val)).otherwise(val)

df = df.withColumn('val_y', apply_another_function(df.val_y))\
       .withColumn('val_x', apply_another_function(df.val_x))

Answer 2

根据我的理解，udf参数是列名。您的示例可能会重写为：

from pyspark.sql.functions import udf, array
from pyspark.sql.types import IntegerType

def change_val_x(val_x):
    threshold = 10
    if val_x > threshold:
        return another_function(val_x)
    else:
        return val_x

def change_val_y(arr):
    threshold = 10
    # arr[0] -> val_x, arr[0] -> val_y 
    if arr[0] > threshold:
        return another_function(arr[1])
    else:
        return val_y

change_val_x_udf = udf(change_val_x, IntegerType())
change_val_y_udf = udf(change_val_y, IntegerType())

# apply these functions to your dataframe
df = df.withColumn('val_y', change_val_y_udf(array('val_x', 'val_y')))\
       .withColumn('val_x', change_val_x_udf('val_x'))

要修改val_x列，一个简单的 udf 就足够了，但对于需要val_y和val_x列值val_y，解决方案是使用数组。请注意，此代码未经测试...

参见此问题对多列应用udf。