如何使用PySpark提取Twitter JSON数据中推文列的用户提及(@)内容

核心思路 #

通过正则匹配识别推文中符合Twitter用户名规则的@内容，直接调用PySpark内置函数或自定义UDF实现提取。
Twitter用户名匹配规则：@开头，后续只能包含大小写字母、数字、下划线，长度为1~15位，对应正则表达式为r'@[a-zA-Z0-9_]{1,15}' #

实现步骤 #

1. 读取数据集 #

首先初始化Spark会话，读取你的Twitter JSON数据集，以下用你提供的示例数据做演示：

from pyspark.sql import SparkSession
from pyspark.sql.functions import regexp_extract_all, col

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

# 真实场景替换为：df = spark.read.json("你的JSON数据集文件路径")
data = [
    ("@ABC", "@XYZ This is my first tweet ever"),
    ("@XYZ", "Hello @ABC welcome to twitter")
]
df = spark.createDataFrame(data, schema=["user", "tweet"])

2. 提取@提及内容 #

Spark 3.1及以上版本可以直接使用内置的regexp_extract_all函数，一次性提取所有匹配的@内容，返回数组类型的结果：

# 匹配@用户的正则规则
mention_regex = r'@[a-zA-Z0-9_]{1,15}'

# 提取所有@提及生成新列mentions
result_df = df.withColumn("mentions", regexp_extract_all(col("tweet"), mention_regex, 0))

# 查看结果
result_df.show(truncate=False)

输出结果如下：

+----+---------------------------------+---------+
|user|tweet                            |mentions |
+----+---------------------------------+---------+
|@ABC|@XYZ This is my first tweet ever |[@XYZ]   |
|@XYZ|Hello @ABC welcome to twitter    |[@ABC]   |
+----+---------------------------------+---------+

3. 低版本Spark兼容方案 #

如果你使用的Spark版本低于3.1，没有regexp_extract_all函数，可以通过自定义UDF实现：

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StringType
import re

# 自定义提取函数
def extract_mentions(text):
    if not text:
        return []
    return re.findall(r'@[a-zA-Z0-9_]{1,15}', text)

# 注册UDF
extract_mentions_udf = udf(extract_mentions, ArrayType(StringType()))

# 调用UDF提取
result_df = df.withColumn("mentions", extract_mentions_udf(col("tweet")))

4. 可选：拆分为单行单提及 #

如果需要把数组中的每个@提及拆成单独的行，配合explode函数即可：

from pyspark.sql.functions import explode

exploded_df = result_df.withColumn("single_mention", explode(col("mentions")))
exploded_df.show(truncate=False)

输出结果：

+----+---------------------------------+---------+--------------+
|user|tweet                            |mentions |single_mention|
+----+---------------------------------+---------+--------------+
|@ABC|@XYZ This is my first tweet ever |[@XYZ]   |@XYZ          |
|@XYZ|Hello @ABC welcome to twitter    |[@ABC]   |@ABC          |
+----+---------------------------------+---------+--------------+

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