熊猫udf分割数组字符串pyspark

Question

我有以下表格

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'

我很困惑，为什么-以及如何修复我的熊猫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

这个问题的答案是

目前，除了MapType、TimestampType的ArrayType和嵌套的StructType之外，所有Spark SQL数据类型都支持基于箭头的转换。

参见此处和此处。