【Python笔记】pyspark dataframe の udf

简单应用

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

mission = "xxx"

spark = SparkSession.builder.appName(mission).enableHiveSupport().getOrCreate()

l =[("A", 16), ("B", 21), ("B", 14), ("B", 18)]
df = spark.createDataFrame(l,["name", "age"])


# 自定义def实现方式1：注册自定义函数
# from pyspark.sql.types import  LongType
def plus_one(a):
    return a + 1

plus_one_udf = udf(plus_one, returnType=LongType())
df.withColumn("one_processed", plus_one_udf(df["age"])).show()

# +----+---+-------------+
# |name|age|one_processed|
# +----+---+-------------+
# |   A| 16|           17|
# |   B| 21|           22|
# |   B| 14|           15|
# |   B| 18|           19|
# +----+---+-------------+

# 自定义def实现方式2：修饰自定义函数
# from pyspark.sql.types import  LongType
@udf(returnType= LongType())
def plus_ten(a):
    return a + 10

df.withColumn("one_processed", plus_ten(df["age"])).show()

# 自定义def实现方式3：Pandas UDFs【最新，最快】
# 但是我们的iterms2里是pyspark2用不了？
from pyspark.sql.functions import pandas_udf, PandasUDFType

#使用 pandas_udf 定义一个 Pandas UDF
#输入/输出都是 double 类型的 pandas.Series
@pandas_udf('long')
def plus_hunderd(a):
    return a + 100
spark.udf.register('plus_hunderd', plus_hunderd)
df.withColumn("one_processed", plus_hunderd(df["age"])).show()

# Grouped Map Pandas UDFs
# 下例显示了使用 grouped map Pandas UDFs从组中的three列每个值减去平均值。每个输入到自定义函数的 pandas.DataFrame 具有相同的 “id” 值。这个用户定义函数的输入和输出模式是相同的，所以我们将“df.schema” 传递给装饰器 pandas_udf 来指定模式。
@pandas_udf(sdf.schema, PandasUDFType.GROUPED_MAP)
#Input/output are both a pandas.DataFrame
def subtract_mean(df):
    return df.assign(new_col=df.three - df.three.mean())

sdf.groupby("id").apply(subtract_mean)

详见: pyspark dataframe之udf

实际场景

def get_group(rate):
    group = (rate*100)//10 + 1
    if (rate*100) % 10 == 0: group -= 1
    if group > 10: group = 10
    return int(group)
    
group_udf = udf(get_group, IntegerType())

df = df1.join(df2,['region','id'],'left')\
        ...
        .withColumn('group',group_udf(col('rate')))