pyspark pandas udf 无法返回任何对象

Question

我正在将我的代码从 Pandas 转移到 Pypark 以执行 NLP 任务。我已经弄清楚如何通过 pandas UDF 应用标记化（使用 Keras 内置库）。但是，我还想返回合适的分词器（供以后在测试数据上使用）。

与 pandas udfs 一样，除了 1-1 列转换（系列、系列列表、缩放器）之外，您不能返回任何其他内容。有什么办法吗？

def tokenize_wrapper(text, maxlen, padding_type):
    tokenizer = Tokenizer(num_words=None, char_level=True, oov_token='UNK')

    @pandas_udf('array<decimal>')
    def tokenize(text):
        tokenizer.fit_on_texts(text)
        names = tokenizer.texts_to_sequences(text)
        padded_data = pad_sequences(names, maxlen=maxlen, padding=padding_type, truncating = padding_type)
        data = np.array(padded_data).tolist()
        return pd.Series(data)
    
    tokenized_names = tokenize(text)
    return tokenized_names