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

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

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.csdn.net/dkjkls