Spark Dataframe lambda 直接在数据帧上

Question

我看到很多例子需要在 rdd.map 上使用 lambda 。
只是想知道我们是否可以做如下的事情：

df.withColumn('newcol',(lambda x: x['col1'] + x['col2'])).show()

Answer 1

您必须将其包装在 UDF 中并提供您希望应用 lambda 的列。

例子：

from pyspark.sql import SparkSession
import pyspark.sql.functions as F

if __name__ == "__main__":
    spark = SparkSession.builder.getOrCreate()
    data = [{"a": 1, "b": 2}]
    df = spark.createDataFrame(data)
    df.withColumn("c", F.udf(lambda x, y: x + y)("a", "b")).show()

结果：

+---+---+---+
|  a|  b|  c|
+---+---+---+
|  1|  2|  3|
+---+---+---+