apache-spark - 在 pyspark[non pandas] 中为数据框的每一行调用一个函数

Question

pyspark中有一个函数:

def sum(a,b):
    c=a+b
    return c

它必须使用 spark sql 在非常非常大的数据帧的每条记录上运行:

x = sum(df.select["NUM1"].first()["NUM1"], df.select["NUM2"].first()["NUM2"])

但这只会对 df 的第一条记录运行它，而不是对所有行运行。我知道这可以使用 lambda 来完成，但我无法以所需的方式对其进行编码。

在现实中； c 将是一个数据框，该函数将执行大量 spark.sql 操作并将其返回。我将不得不为每一行调用该函数。我想，我会尝试使用这个 sum(a,b) 作为类比来理解它。

+----------+----------+-----------+
|     NUM1 |     NUM2 |    XYZ    |
+----------+----------+-----------+
|      10  |     20   |      HELLO|                                    
|      90  |     60   |      WORLD|
|      50  |     45   |      SPARK|
+----------+----------+-----------+


+----------+----------+-----------+------+
|     NUM1 |     NUM2 |    XYZ    | VALUE|
+----------+----------+-----------+------+
|      10  |     20   |      HELLO|30    |                                     
|      90  |     60   |      WORLD|150   |
|      50  |     45   |      SPARK|95    |
+----------+----------+-----------+------+

Python: 3.7.4
Spark: 2.2

Answer 1

您可以使用 .withColumn 函数:

from pyspark.sql.functions import col
from pyspark.sql.types import LongType
df.show()
+----+----+-----+
|NUM1|NUM2|  XYZ|
+----+----+-----+
|  10|  20|HELLO|
|  90|  60|WORLD|
|  50|  45|SPARK|
+----+----+-----+

def mysum(a,b):
  return a + b

spark.udf.register("mysumudf", mysum, LongType())

df2 = df.withColumn("VALUE", mysum(col("NUM1"),col("NUM2"))

df2.show()
+----+----+-----+-----+
|NUM1|NUM2|  XYZ|VALUE|
+----+----+-----+-----+
|  10|  20|HELLO|   30|
|  90|  60|WORLD|  150|
|  50|  45|SPARK|   95|
+----+----+-----+-----+