PySpark pandas_udf数据类型错误：无法将str转为int32求助

问题场景 #

实现position_words普通UDF和filter_words Pandas UDF处理文本数据时，抛出类型转换错误：

pyarrow.lib.ArrowInvalid: Could not convert 'position' with type str: tried to convert to int32

原UDF代码 #

@udf_annotator(returnType=ArrayType(StructType([StructField("position", IntegerType(), True),
                                       StructField("tokens", StringType(), True)])))
def position_words(tokens):
    position = [(int(i), token) for i, token in enumerate(tokens)]
    return position

@pandas_udf(returnType=ArrayType(StructType([StructField("position", IntegerType(), True),
                                                StructField("word", StringType(), True)])))
def filter_words(lst2, lang2):
    def filter_word2(lst, lang):
        filtered_tokens = []
        for pos, word in lst:
            if word is None: continue
            if len(word) == 0: continue
            text = re.sub(
                r"((https?|ftps?|file)?:\/\/)?(?:[\w\d!#$&'()*\+,;:?@[\]\-_.~]|(?:%[0-9a-fA-F][0-9a-fA-F]))+" +
                "\\.([\w\d]{2,6})(\/(?:[\w\d!#$&'()*\+,;:?@[\]\-_.~]|(?:%[0-9a-fA-F][0-9a-fA-F]))+)*",
                "", word)
            text = re.sub(r"[@#]\w+", "", text)
            text = re.sub(r"'", " ", text)
            word_filtered = re.findall(r"""(?:\w\.)+\w\.?|\w{2,}-?\w*-*\w*""", text)
            word_filtered = " ".join(word_filtered)
            filtered_tokens.append((pos, word_filtered))
        return filtered_tokens
    all_founded_result = [filter_word2(lst, lang) for lst, lang in zip(lst2, lang2)]
    return pd.Series(all_founded_result)

测试代码 #

import random
langs = ['eng', 'rus', 'tuk', 'jpn', 'arb', 'fin', 'fra', 'cmn']

def random_text(length):
    return ''.join(random.choice('sdfsdfg jkhkhkj jh kh') for _ in range(length))

df = pd.DataFrame({'text': [random_text(10) for _ in range(100000)],
                       'lang': [random.choice(langs) for _ in range(100000)]})
sdf = spark.createDataFrame(df).withColumn('tokens', F.split(F.col('text')))\
  .withColumn("position", position_words(F.col("tokens")))\
  .withColumn("position_filt", filter_words(F.col("position"), F.col("lang")))

错误原因 #

普通UDF返回的结构体数组通过Arrow序列化传递给Pandas UDF时，position字段被解析为字符串类型（与定义的IntegerType不符），但Pandas UDF返回类型要求该字段为整数，直接使用未转换的pos导致类型不匹配。

解决方案 #

在filter_word2函数的循环中，显式将pos转换为整数，确保类型与返回定义一致。

修正后的代码 #

import re
import pandas as pd
from pyspark.sql.functions import pandas_udf, udf, split as F_split
from pyspark.sql.types import ArrayType, StructType, StructField, IntegerType, StringType

# 若udf_annotator为自定义装饰器，可替换为普通@udf
@udf(returnType=ArrayType(StructType([StructField("position", IntegerType(), True),
                                       StructField("tokens", StringType(), True)])))
def position_words(tokens):
    position = [(int(i), token) for i, token in enumerate(tokens)]
    return position

@pandas_udf(returnType=ArrayType(StructType([StructField("position", IntegerType(), True),
                                                StructField("word", StringType(), True)])))
def filter_words(lst2, lang2):
    def filter_word2(lst, lang):
        filtered_tokens = []
        for pos, word in lst:
            if word is None: continue
            if len(word) == 0: continue
            # 显式转换pos为整数
            pos_int = int(pos)
            text = re.sub(
                r"((https?|ftps?|file)?:\/\/)?(?:[\w\d!#$&'()*\+,;:?@[\]\-_.~]|(?:%[0-9a-fA-F][0-9a-fA-F]))+" +
                "\\.([\w\d]{2,6})(\/(?:[\w\d!#$&'()*\+,;:?@[\]\-_.~]|(?:%[0-9a-fA-F][0-9a-fA-F]))+)*",
                "", word)
            text = re.sub(r"[@#]\w+", "", text)
            text = re.sub(r"'", " ", text)
            word_filtered = re.findall(r"""(?:\w\.)+\w\.?|\w{2,}-?\w*-*\w*""", text)
            word_filtered = " ".join(word_filtered)
            filtered_tokens.append((pos_int, word_filtered))
        return filtered_tokens
    all_founded_result = [filter_word2(lst, lang) for lst, lang in zip(lst2, lang2)]
    return pd.Series(all_founded_result)

验证说明 #

使用原测试代码调用修正后的UDF，将不再出现类型转换错误，position_filt字段会正确返回包含整数位置和过滤后文本的结构体数组。

内容的提问来源于stack exchange，提问作者Rory