Spark NLP与PySpark结合构建模型报错：无法识别函数类型Pipeline阶段

错误原因 #

你遇到的TypeError是因为直接将UDF函数添加到了Pipeline的stages中，Spark Pipeline的每个阶段必须是Transformer或Estimator类型的实例，而UDF是普通函数，无法被Pipeline识别。

解决方案 #

不需要自定义UDF转换嵌入向量，Spark NLP提供了EmbeddingsFinisher组件，可以直接将UniversalSentenceEncoder输出的嵌入注解转换成PySpark的DenseVector类型，完美适配Pipeline流程。同时修正Pipeline阶段的逻辑，确保所有阶段都是合法的Transformer/Estimator。

修改后的完整代码 #

from pyspark.sql.functions import col
from pyspark.ml import Pipeline
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler
from sparknlp.base import DocumentAssembler, EmbeddingsFinisher
from sparknlp.annotator import UniversalSentenceEncoder

# 定义分类列
categorical_cols = ["keyword", "location"]

# 构建独热编码阶段
stages_one_hot = []
for col_name in categorical_cols:
    indexer = StringIndexer(inputCol=col_name, outputCol=f"{col_name}_index")
    encoder = OneHotEncoder(inputCol=f"{col_name}_index", outputCol=f"{col_name}_vec")
    stages_one_hot.extend([indexer, encoder])

# 构建文本嵌入阶段
stages_text = []
# 文本转文档注解
documentAssembler = DocumentAssembler().setInputCol("text").setOutputCol("document")
# 通用句子编码器
use_encoder = UniversalSentenceEncoder.pretrained().setInputCols(["document"]).setOutputCol("sentence_embedding")
# 嵌入向量转换为DenseVector
embeddings_finisher = EmbeddingsFinisher() \
    .setInputCols(["sentence_embedding"]) \
    .setOutputCols(["sentence_embedding_dense"]) \
    .setOutputAsVector(True)  # 输出为PySpark的Vector类型

stages_text.extend([documentAssembler, use_encoder, embeddings_finisher])

# 构建特征组装与模型阶段
stages_together = []
# 合并分类特征与文本嵌入特征
assembler = VectorAssembler(
    inputCols=[f"{col}_vec" for col in categorical_cols] + ["sentence_embedding_dense"],
    outputCol="features"
)
# 逻辑回归模型
lr = LogisticRegression(featuresCol="features", labelCol="target")

stages_together.extend([assembler, lr])

# 构建完整Pipeline
pipeline = Pipeline(stages=stages_one_hot + stages_text + stages_together)

# 训练与转换
df = trainDataset
pipelineModel = pipeline.fit(df)
transformed_df = pipelineModel.transform(df)

关键修改说明 #

• 移除了自定义UDF，改用EmbeddingsFinisher：该组件是Spark NLP官方提供的Transformer，完全符合Pipeline要求，能直接将嵌入注解转为DenseVector，避免类型错误。
• 修正了变量名冲突：原代码中用encoder同时指代OneHotEncoder和UniversalSentenceEncoder，修改为use_encoder避免混淆。
• 调整了阶段添加方式：用extend替代+=，逻辑更清晰。

你完全可以完成这种特征提取并构建模型，Spark NLP和PySpark的Pipeline天然兼容，只要确保所有Pipeline阶段都是合法的Transformer/Estimator实例即可。

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