PySpark DataFrame按行提取高频词并按规则输出的实现需求

现有如下结构的PySpark DataFrame：

columns = ["object_type", "object_name"]
data = [("galaxy", "andromeda,milky way,condor,andromeda"),
        ("planet", "mars,jupiter,venus,mars,saturn,venus,earth,mars,venus,earth"), 
        ("star", "mira,sun,altair,sun,sirius,rigel,mira,sirius,aldebaran"),
        ("natural satellites", "moon,io,io,elara,moon,kale,titan,kale,phobos,titan,europa")]
init_df = spark.createDataFrame(data).toDF(*columns)
init_df.show(truncate = False)

输出结果：

+------------------+-----------------------------------------------------------+
|object_type       |object_name                                                |
+------------------+-----------------------------------------------------------+
|galaxy            |andromeda,milky way,condor,andromeda                       |
|planet            |mars,jupiter,venus,mars,saturn,venus,earth,mars,venus,earth|
|star              |mira,sun,altair,sun,sirius,rigel,mira,sirius,aldebaran     |
|natural satellites|moon,io,io,elara,moon,kale,titan,kale,phobos,titan,europa  |
+------------------+-----------------------------------------------------------+

需要新增most_frequent列，提取object_name中的高频词，规则如下：

• 仅1个主导词（众数数量为1）：取该词（如第一行的"andromeda"）
• 2个出现次数相同的主导词：同时取这两个词（如第二行的"mars"和"venus"）
• 3个出现次数相同的主导词：取这三个词（如第三行的"mira"、"sun"和"sirius"）
• 4个及以上出现次数相同的主导词：设置为"many objects"（如第四行）

预期输出：

+------------------+-----------------------------------------------------------+---------------+
|object_type       |object_name                                                |most_frequent  |
+------------------+-----------------------------------------------------------+---------------+
|galaxy            |andromeda,milky way,condor,andromeda                       |andromeda      |
|planet            |mars,jupiter,venus,mars,saturn,venus,earth,mars,venus,earth|mars,venus     |
|star              |mira,sun,altair,sun,sirius,rigel,mira,sirius,aldebaran     |mira,sun,sirius|
|natural satellites|moon,io,io,elara,moon,kale,titan,kale,phobos,titan,europa  |many objects   |
+------------------+-----------------------------------------------------------+---------------+

以下是完整的PySpark代码实现，分步骤处理：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 1. 拆分object_name为单个元素的数组，并展开成多行
exploded_df = init_df.withColumn("obj", F.split(F.col("object_name"), ",")) \
    .withColumn("obj", F.explode(F.col("obj")))

# 2. 按object_type和obj分组，统计每个object的出现次数
count_df = exploded_df.groupBy("object_type", "obj") \
    .agg(F.count("obj").alias("count"))

# 3. 定义窗口：按object_type分组，按count降序排序
window_spec = Window.partitionBy("object_type").orderBy(F.desc("count"))

# 4. 给每个object_type下的obj标记排名，同时获取当前组的最大count值
ranked_df = count_df.withColumn("rank", F.rank().over(window_spec)) \
    .withColumn("max_count", F.max("count").over(Window.partitionBy("object_type")))

# 5. 筛选出等于最大count的obj，按object_type分组收集这些obj，并统计数量
grouped_df = ranked_df.filter(F.col("count") == F.col("max_count")) \
    .groupBy("object_type") \
    .agg(
        F.collect_list("obj").alias("top_objs"),
        F.count("obj").alias("top_count")
    )

# 6. 根据top_count的值生成most_frequent列
result_df = init_df.join(grouped_df, on="object_type", how="inner") \
    .withColumn("most_frequent", 
                F.when(F.col("top_count") >=4, "many objects")
                .otherwise(F.concat_ws(",", F.col("top_objs")))) \
    .drop("top_objs", "top_count")

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

步骤解释 #

1. 拆分与展开：用split把逗号分隔的字符串转成数组，再用explode将数组展开为多行，方便后续统计次数。
2. 统计次数：按object_type和单个obj分组，统计每个对象的出现次数。
3. 排名与最大次数：通过窗口函数对每个object_type下的对象按次数降序排名，同时获取当前组的最大出现次数，筛选出所有等于最大次数的对象。
4. 收集与判断：将同一object_type下的高频对象收集为列表，统计高频对象的数量，再根据数量规则生成最终的most_frequent列。
5. 合并结果：将处理后的结果与原DataFrame关联，去掉中间列，得到最终输出。

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