Pyspark：访问UDF中行中的列

Question

pyspark的初学者试图了解UDF：

我有一个PySpark数据帧p_b，我通过传递数据帧的所有行来调用UDF。我想从行访问列debit。由于某种原因，这没有发生。请在下面找到代码段。

p_b has 4 columns, id, credit, debit,sum

功能：

def test(row): return('123'+row['debit'])

转换为UDF

test_udf=udf(test,IntegerType())

在数据帧p_b上调用UDF

vals=test_udf(struct([p_b[x] for x in p_b.columns])) print(type(vals)) print(vals)

输出

Column<b'test(named_struct(id, credit,debit,sum))'>

Answer 1

首先创建一个示例数据框：

from pyspark.sql.functions import *
from pyspark.sql.types import *  
schema = StructType([StructField("id", StringType(), True),\
                               StructField("credit", IntegerType(), True),\
                       StructField("debit", IntegerType(), True),\
                     StructField("sum", IntegerType(), True)])
df = spark.createDataFrame([("user_10",100, 10,110),("user_11",200, 20,220),("user_12",300, 30,330) ], schema)
df.show()

结果为：

+-------+------+-----+---+
|     id|credit|debit|sum|
+-------+------+-----+---+
|user_10|   100|   10|110|
|user_11|   200|   20|220|
|user_12|   300|   30|330|
+-------+------+-----+---+

现在，让我们定义将123添加到传递给它的值上的udf：

def test(x):
    return(123+x)
test_udf=udf(test,IntegerType())

让我们看看如何使用UDF：

df2 = df.withColumn( 'debit' , test_udf(col('debit')) )
df2.show()

结果为：

+-------+------+-----+---+
|     id|credit|debit|sum|
+-------+------+-----+---+
|user_10|   100|  133|110|
|user_11|   200|  143|220|
|user_12|   300|  153|330|
+-------+------+-----+---+

请注意，现在您可能需要重新计算“求和”列：

df2 = df2.withColumn( 'sum' ,  col('debit')+col('credit') )
df2.show()

结果为：

+-------+------+-----+---+
|     id|credit|debit|sum|
+-------+------+-----+---+
|user_10|   100|  133|233|
|user_11|   200|  143|343|
|user_12|   300|  153|453|
+-------+------+-----+---+