pandas udf 拆分字符串数组 pyspark

Question

我有以下table

id | country_mapping
--------------------
1  | {"GBR/bla": 1,
      "USA/bla": 2}

我想创建一个包含以下内容的列

id | source_countries
--------------------
1  | ["GBR", "USA"]

我需要通过 pandas udf 来完成。我创建了以下内容

import pyspark.sql.functions as F

@F.pandas_udf("string")
def func(s):
    return s.apply(lambda x: [y.split("/")[0] for y in x])

我认为这行得通，因为如果我 运行 纯 pandas 这段代码，它可以满足我的需要。

import pandas as pd
s = pd.Series([["GBR/1", "USA/2"], ["ITA/1", "FRA/2"]])
s.apply(lambda x: [y.split("/")[0] for y in x])

给予

Out[1]: 0    [GBR, USA]
        1    [ITA, FRA]
dtype: object

但是当我运行

df.withColumn('source_countries', 
              func(F.map_keys(F.col("country_mapping")))).collect()

当我 运行 以下内容时失败并出现以下错误：

PythonException: An exception was thrown from a UDF: 'pyarrow.lib.ArrowTypeError: Expected bytes, got a 'list' object'

我很困惑为什么 - 以及如何修复我的 pandas udf。

Answer 1

而不是 pandas_udf，您可以以类似的方式使用 udf

from pyspark.sql import functions as F
from pyspark.sql import types as T

def func(v):
    return [x.split('/')[0] for x in v]

(df
     .withColumn('source_countries', F.udf(func, T.ArrayType(T.StringType()))(F.map_keys(F.col('country_mapping'))))
     .show(10, False)
)

# +---+----------------------------+----------------+
# |id |country_mapping             |source_countries|
# +---+----------------------------+----------------+
# |1  |{USA/bla -> 2, GBR/bla -> 1}|[USA, GBR]      |
# +---+----------------------------+----------------+

Answer 2

这个问题的答案是

Currently, all Spark SQL data types are supported by Arrow-based conversion except MapType, ArrayType of TimestampType, and nested StructType.

见here and here。