PySpark CheatSheet-UDF相关

1. PySpark的各种输入输出UDF：

下面展示了多个例子，包含：

0. UDF单列输入，单列输出的简化版，只用一行代码（lambda函数）；

1. UDF单列输入，单列输出；

2. UDF多列输入，单列输出；

3. UDF单列输入，多列输出；

4. UDF多列输入，多列输出；

5. UDF单列输入并添加外置参数，单列输出。

6. UDF多列输入并添加外置参数，单列输出。

只有第一个UDF括号里头的xxxType()可以删除。所以对于稍微复杂一些的UDF，就要给udf()函数设计输出格式，以后如果还要设计输入格式，后续会再修改。

另外，看看最后的df4、df5、df6（4、5一样，6多了一列）是什么样子的，也是可以了解到“*”的使用技巧。

from pyspark.sql.functions import udf, struct
from pyspark.sql.types import *

# 建立测试DataFrame
df = sc.parallelize([(1, 2.0), (2, 3.0)]).toDF(["x", "y"])

# UDF单列输入，单列输出简化版，只用一行lambda
df0 = df.withColumn("half_x", udf(lambda x: 0.5 * x)('x'))

# UDF单列输入，单列输出
def half(n):
    return n / 2 

udf_half = udf(half, FloatType())

df1 = df.select("x", "y", udf_half("y").alias("half_y"))

# UDF多列输入，单列输出
def sum(x,y):
    return x + y

udf_sum = udf(sum, FloatType())

df2 = df.select("x", "y", udf_sum("x","y").alias("x+y"))

# UDF单列输入，多列输出
def dou_tri(x):
    return 2*x, 3.0*x

schema = StructType([
    StructField("double", IntegerType(), False),
    StructField("triple", FloatType(), False)
])
udf_dou_tri = udf(dou_tri, schema)

df3 = df.select("x", "y", udf_dou_tri("x").alias("dou_tri")).select("x", "y", "dou_tri.*")

# UDF多列输入，多列输出
def comp(x,y):
    return x+y, x-y

schema = StructType([
    StructField("sum", FloatType(), False),
    StructField("diff", FloatType(), False)
])
udf_comp = udf(comp, schema)
df4 = df.select("x", "y", udf_comp("x", "y").alias("sum_dif")).select("x", "y", "sum_dif.*")
df5 = df.select("*", udf_comp("x", "y").alias("sum_dif")).select("x", "y", "sum_dif.*")
df6 = df.select("*", udf_comp("x", "y").alias("sum_dif")).select("*", "sum_dif.*").drop("sum_dif")

# UDF单列输入并添加外置参数，单列输出
def sum(x,y): 
    return x + y

def udf_sum_y(y): 
    return udf(lambda x: sum(x,y))

df7 = df.withColumn("x+y", udf_sum_y(100)("x"))

# UDF多列输入并添加外置参数，单列输出
def sum(x, y, z): 
    return x + y + z

def udf_sum_z(z): 
    return udf(lambda x, y: sum(x,y, z))

df8 = df.withColumn("x+y+z", udf_sum_z(100)("x", "y"))