如何用PySpark实现保留多词实体的文本分词？

要实现不拆分指定多词游戏名称的分词，核心思路是优先匹配长实体，再处理普通单词，以下是两种可行方案：

方案一：使用RegexTokenizer（快速实现） #

利用正则表达式优先匹配游戏实体列表中的长名称，再匹配其他单个单词，通过RegexTokenizer直接生成目标分词结果。

代码示例 #

from pyspark.sql import SparkSession
from pyspark.ml.feature import RegexTokenizer
import re

# 初始化Spark会话
spark = SparkSession.builder.appName("GameNERTokenizer").getOrCreate()

# 测试数据
data = [("Dead Space Remake Is Only $28 At GameFly.",)]
df = spark.createDataFrame(data, ["text"])

# 游戏实体列表
games = ["Dead Space", "Dead Space Remake", "Witcher 3"]

# 关键步骤：按实体长度降序排序，优先匹配长实体避免拆分
sorted_games = sorted(games, key=lambda x: len(x), reverse=True)
# 转义实体中的特殊字符，避免正则匹配异常
game_patterns = [re.escape(game) for game in sorted_games]
# 构建正则：先匹配游戏实体，再匹配任意非空白字符序列（即普通单词/符号）
regex_pattern = f"({'|'.join(game_patterns)})|\\S+"

# 初始化Tokenizer，gaps=False表示匹配pattern内容作为分词
tokenizer = RegexTokenizer(
    inputCol="text",
    outputCol="tokens",
    pattern=f"(?!{regex_pattern})",
    gaps=False
)

# 执行分词
tokenized_df = tokenizer.transform(df)
tokenized_df.select("tokens").show(truncate=False)

输出结果 #

+------------------------------------------------+
|tokens                                          |
+------------------------------------------------+
|[Dead Space Remake, Is, Only, $28, At, GameFly.]|
+------------------------------------------------+

如果需要忽略大小写匹配，可在正则前添加(?i)：

regex_pattern = f"(?i)({'|'.join(game_patterns)})|\\S+"

方案二：自定义Transformer（灵活扩展） #

如果需要动态更新实体列表或处理更复杂的匹配逻辑，可以自定义Spark Transformer，封装分词逻辑。

代码示例 #

from pyspark.sql import SparkSession
from pyspark.ml import Transformer
from pyspark.ml.param.shared import HasInputCol, HasOutputCol, Param
from pyspark.sql.types import ArrayType, StringType
from pyspark.sql.functions import udf, col
import re

spark = SparkSession.builder.appName("CustomGameTokenizer").getOrCreate()

# 自定义实体保留分词器
class EntityPreservingTokenizer(Transformer, HasInputCol, HasOutputCol):
    def __init__(self, inputCol=None, outputCol=None, entities=None):
        super().__init__()
        self.entities = Param(self, "entities", "List of multi-word entities to preserve")
        self._setDefault(entities=[])
        self.setParams(inputCol=inputCol, outputCol=outputCol, entities=entities)
    
    def setParams(self, inputCol=None, outputCol=None, entities=None):
        return self._set(inputCol=inputCol, outputCol=outputCol, entities=entities)
    
    def _transform(self, df):
        entities = self.getOrDefault(self.entities)
        # 按长度降序排序实体
        sorted_entities = sorted(entities, key=lambda x: len(x), reverse=True)
        entity_patterns = [re.escape(e) for e in sorted_entities]
        pattern = f"({'|'.join(entity_patterns)})|\\S+"
        
        # 定义分词UDF
        def tokenize(text):
            return re.findall(pattern, text)
        
        tokenize_udf = udf(tokenize, ArrayType(StringType()))
        return df.withColumn(self.getOutputCol(), tokenize_udf(col(self.getInputCol())))

# 使用自定义分词器
data = [("Dead Space Remake Is Only $28 At GameFly.",)]
df = spark.createDataFrame(data, ["text"])
games = ["Dead Space", "Dead Space Remake", "Witcher 3"]

custom_tokenizer = EntityPreservingTokenizer(
    inputCol="text",
    outputCol="tokens",
    entities=games
)

tokenized_df = custom_tokenizer.transform(df)
tokenized_df.select("tokens").show(truncate=False)

核心逻辑说明 #

• 必须按实体长度降序排序：确保长实体（如Dead Space Remake）先被匹配，避免被短实体（如Dead Space）拆分。
• 正则转义：用re.escape()处理实体中的特殊字符，避免正则语法冲突。

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