Pyspark 使用 Spark Udf 的一些经验

起初开始写一些 udf 的时候感觉有一些奇怪，在 spark 的计算中，一般通过转换(Transformation) 在不触发计算(Action) 的情况下就行一些预处理。udf 就是这样一个好用的东西，他可以在我们进行 Transformation 的时候给我们带来对复杂问题的处理能力。

这里有两种最典型的方法。

应用于 spark 2.4

1. 直接在 SparkSession.sql 里面直接使用注册好的 udf，类似于这种写法

xx = SparkSession.catalog.registerFunction('fmt_buy_channel', lambda i, j, x, y: HdNewOrderRecord.fmt_buy_channel(i, j, x, y))

ss.sql("""
           SELECT t1.pay_id,
           t1.sku_mode,
           LEFT(t1.charge_time, 19) AS buy_time,
           fmt_buy_channel(t1.join_type, t1.special_card_type, t1.channel_type, t1.pay_channel) AS channel,
           t1.pay_money,
           t1.charge_user_id
           FROM analytics_db.hd_new_order_record t1 JOIN user_info t2
           ON (t1.charge_user_id = t2.user_id
           AND t1.charge_time < '{}') ORDER BY t1.charge_time ASC
       """.format(dump_time))

可以看到我们定义的 udf "fmt_buy_channel" 被直接用在了 sql 语句里面。这种 spark 是可以轻松处理的。不过这种写法有个问题，在使用了 udf 之后，这个字段不能立即嵌套另外的 function 。否则可能会报错，比如我写一个这样的函数

df = ss.sql("""
                SELECT t1.pay_id,
                t1.sku_mode,
                LEFT(t1.charge_time, 19) AS buy_time,
                fmt_buy_channel(t1.join_type, t1.special_card_type, t1.channel_type, t1.pay_channel) AS channel,
                t1.pay_money,
                t1.charge_user_id
                FROM analytics_db.hd_new_order_record t1 JOIN user_info t2
                ON (t1.charge_user_id = t2.user_id
                AND t1.charge_time < '{}') ORDER BY t1.charge_time ASC
            """.format(dump_time))

会无法正确执行。

2. 第二种方法是我们可以直接使用 pyspark 提供的函数进行 udf 调用，pyspark 或者本身的 scala spark 他们为我们封装了非常多基于 SparkSession 和 DataFrame 的函数。

来看一个结合了两者的一个完整的例子

df = ss.sql("""
                SELECT t1.pay_id,
                t1.sku_mode,
                LEFT(t1.charge_time, 19) AS buy_time,
                fmt_buy_channel(t1.join_type, t1.special_card_type, t1.channel_type, t1.pay_channel) AS channel,
                t1.pay_money,
                t1.charge_user_id
                FROM analytics_db.hd_new_order_record t1 JOIN user_info t2
                ON (t1.charge_user_id = t2.user_id
                AND t1.charge_time < '{}') ORDER BY t1.charge_time ASC
            """.format(dump_time))
df = df.select(df.charge_user_id, concat_ws('_', df.pay_id, df.channel, df.sku_mode, df.buy_time, df.pay_money).alias('sku_buys'))\
    .groupBy(df.charge_user_id)\
    .agg(collect_list('sku_buys').alias('sku_buys'))
df.createOrReplaceTempView(table_name)

上面我使用了常用的一些 SQL 函数，其实 spark 对这些函数都有包装 。比如 left 之类的函数都可以在 pyspark.sql.functions import 中找到例如 ltrim。

第一条语句我们通过 ss.sql 获得一个 df 。

第二条语句我们通过操纵 df 的函数生成我们自己需要的字段，并且对字符串进行拼接。最后分组展示。这里用到了几个函数需要介绍一下。

concat_ws: concat_ws 用于拼接字符串，第一个参数接受一个拼接用的符号，后面依次跟上需要拼接的字段即可。

.groupBy().agg(collect_list): 在被基于某一项分组之后，可以使用 spark 提供的 agg 来接收一个聚合函数。 collect_list 这里可以将分组的多个字段基于被 group by 的字段拼接成一个 list 。他还有一个类似功能的函数是 collect_set，在拼接的时候会去重被 append 的数据。

新老版本 spark 在 udf 的使用上会有一些位置上的不一样。特别是在 1.6 跨度到 2.0 的时候。之前还看到过另外一个注册使用方法，放出来给大家看。

from pyspark.sql.functions import udf
from pyspark.sql.types import BooleanType

def regex_filter(x):
    regexs = ['.*ALLYOURBASEBELONGTOUS.*']

    if x and x.strip():
        for r in regexs:
            if re.match(r, x, re.IGNORECASE):
                return True

    return False 

filter_udf = udf(regex_filter, BooleanType())

df_filtered = df.filter(filter_udf(df.field_to_filter_on))

这个跟上面的注册方法最终都会走到 udf 的注册和 udf._wrapped 这个方法并且返回一个函数。如果不接收这个函数返回值，那么可以直接在 ss.sql 中当 udf 进行使用。如果接收当函数值，可以放在 df 的函数里面方便的进行使用。

另外在 spark 2.4 版本以前的 2.2 版本，要想直接获得一个注册完毕的 udf 不能使用上面的 register 方法。那个方法在 2.3 追加 return 。如果我们需要 return 一个 udf 对象我们要这样做

import pyspark.sql.functions as f
right_user = f.udf(lambda i, j, x, y, o, p: HdNewUserInfo.right_user(i, j, x, y, o, p))

使用 udf + sql 函数可以方便的帮助我们进行 transformation ，来完成更加复杂的的计算逻辑。

Reference:

https://stackoverflow.com/questions/31816975/how-to-pass-whole-row-to-udf-spark-dataframe-filter   How to pass whole Row to UDF - Spark DataFrame filter

https://stackoverflow.com/questions/52051985/filter-pyspark-dataframe-with-udf-on-entire-row/52055861   Filter Pyspark Dataframe with udf on entire row

https://gist.github.com/samuelsmal/feb86d4bdd9a658c122a706f26ba7e1e   pyspark_udf_filtering.py

https://stackoverflow.com/questions/36784000/how-to-filter-a-spark-dataframe-by-a-boolean-column   how to filter a spark dataframe by a boolean column

https://stackoverflow.com/questions/37580782/pyspark-collect-set-or-collect-list-with-groupby   pyspark collect_set or collect_list with groupby

https://www.jianshu.com/p/bded081b5350

https://www.cnblogs.com/fudashi/p/7491039.html

https://gist.github.com/samuelsmal/feb86d4bdd9a658c122a706f26ba7e1e