SparkSQL UDF 返回类型不能包含NUMPY类型

最新推荐文章于 2023-10-31 17:54:12 发布

dkjkls

最新推荐文章于 2023-10-31 17:54:12 发布

阅读量1.1k

点赞数 1

CC 4.0 BY-SA版权

分类专栏：大数据文章标签： spark udf numpy pandas 大数据

本文链接：https://blog.youkuaiyun.com/dkjkls/article/details/90742439

大数据专栏收录该内容

5 篇文章

订阅专栏

本文探讨了SparkSQL中用户自定义函数(UDF)在处理包含Numpy类型数据时遇到的PickleException错误。通过实例演示了错误产生的原因及解决方案，强调了在Spark2.1.1版本中，UDF不支持直接使用Numpy类型，并提供了代码修正方法。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

版权声明：本文为博主原创文章，未经博主允许不得转载。

Spark版本: 2.1.1

问题：SparkSQL UDF 返回类型包含NUMPY类型，报错：

PickleException: expected zero arguments for construction of ClassDict (for numpy.dtype)

如下方代码所示，定义的udf返回类型为包含dict的list，在udf中新增排序列并赋默认值，如此写会报上述错误。

import pandas as pd
from pyspark.sql import functions as F
from pyspark.sql.types import MapType, StringType, ArrayType


def calc_allocate(collect_list):
    df_assign = pd.DataFrame(collect_list)
    # 新增排序列并赋默认值
    df_assign["rank"] = range(1, df_assign.shape[0] + 1)
    return_allocate = []

    for i in range(1, df_assign.shape[0] + 1):
        allocate_dict = dict()
        allocate_dict["id"] = df_assign.loc[df_assign["rank"] == i, "id"].values[0]
        allocate_dict["rank"] = df_assign.loc[df_assign["rank"] == i, "rank"].values[0]
        return_allocate.append(allocate_dict)

    return return_allocate


if __name__ == '__main__':
    collect_list = [
        {u'id': u'1'},
        {u'id': u'2'},
        {u'id': u'3'},
        {u'id': u'4'},
        {u'id': u'5'}
    ]
    data_allocate_compute = calc_allocate(collect_list)
    # 定义udf
    udf_calc_allocate = F.udf(collect_list, returnType=ArrayType(MapType(StringType(), StringType())))

参见Spark2.1.1的官方文档，可知SparkSQL不能使用引用的类型，如numpy中的数据类型。

Python DataTypes No Longer Singletons

When using DataTypes in Python you will need to construct them (i.e. StringType()) instead of referencing a singleton.

——摘自http://spark.apache.org/docs/2.1.1/sql-programming-guide.html#python-datatypes-no-longer-singletons

查看字段类型可知，在新增排序列“rank”时，类型为 numpy.int64 ，所以会保报错。

print type(df_assign.loc[df_assign["rank"] == i, "rank"].values[0])
<type 'numpy.int64'>

解决方式：在赋值时强制转换python基础类型。

allocate_dict["rank"] = int(df_assign.loc[df_assign["rank"] == i, "rank"].values[0])

而在高版本中会支持numpy返回新型，如2.4.3的文档。

http://spark.apache.org/docs/latest/sql-pyspark-pandas-with-arrow.html#pandas-udfs-aka-vectorized-udfs

--------------------------文档信息--------------------------
版权声明：本文为博主原创文章，未经博主允许不得转载
署名(BY) ：dkjkls（dkj卡洛斯）
文章出处：http://blog.youkuaiyun.com/dkjkls