PySpark系列：pyspark udf的使用

PySpark系列：udf的使用

前言

pyspark.sql.functions提供了很多预定义的函数用来对列数据进行处理，有三角函数、数学函数、agg相关函数、窗口函数、字符串处理函数、列编解码函数、时间相关函数等。但在实际使用中，依然会遇到很多复杂的数据结构，下面举例来说明udf的使用。
官网上对udf的说明：pyspark.sql.functions.udf.

本文使用的一些包和模块:

import pandas as pd
from pyspark.sql.types import IntegerType
from pyspark.sql import functions as F
from pyspark.sql import SparkSession
spark=SparkSession.builder.appName("test").enableHiveSupport().getOrCreate()

1. 直接使用udf

比如我们有如下的df：

d1 = {'name':["赵四","刘能","广坤","浩哥"], 
      'height_weight':["165,70","163,75","167,65","176,60"]}
df1 = spark.createDataFrame(pd.DataFrame(d1))
df1.show()
"""
输出:
+----+-------------+
|name|height_weight|
+----+-------------+
| 赵四|       165,70|
| 刘能|       163,75|
| 广坤|       167,65|
| 浩哥|       176,60|
+----+-------------+
"""

我们要把用户的体重拿出来，并设置为Integer。

# 定义udf
get_weight = F.udf(lambda s: int(s.split(",")[1]), IntegerType())
df2 = df1.withColumn("weight", get_weight(F.col("height_weight")))
df2.show()
"""
输出:
+----+-------------+------+
|name|height_weight|weight|
+----+-------------+------+
| 赵四|       165,70|    70|
| 刘能|       163,75|    75|
| 广坤|       167,65|    65|
| 浩哥|       176,60|    60|
+----+-------------+------+
"""

说明：

1. 第2行代码的int必须要加上，否则df2.weight为null；
2. 第2行代码的IntegerType()必须要加上，否则即使加上int，df2.weight类型依旧为string；

2. udf作为不带参装饰器

不能处理udf的返回类型，只能是string

@F.udf
def get_weight(s):
    return s.split(",")[1]
df2 = df1.withColumn("weight", get_weight(F.col("height_weight")))
df2.show()
# 其中df2.weight是string类型
"""
输出:
+----+-------------+------+
|name|height_weight|weight|
+----+-------------+------+
| 赵四|       165,70|    70|
| 刘能|       163,75|    75|
| 广坤|       167,65|    65|
| 浩哥|       176,60|    60|
+----+-------------+------+
"""

3. udf作为带参装饰器

可以处理udf的返回类型

@F.udf(returnType=IntegerType())
def get_weight(s):
    return int(s.split(",")[1])
df2 = df1.withColumn("weight", get_weight(F.col("height_weight")))
df2.show()
# 其中df2.weight是int类型
"""
输出:
+----+-------------+------+
|name|height_weight|weight|
+----+-------------+------+
| 赵四|       165,70|    70|
| 刘能|       163,75|    75|
| 广坤|       167,65|    65|
| 浩哥|       176,60|    60|
+----+-------------+------+
"""

完！