请求：将Python分组过滤代码转为PySpark实现，支持百万级数据处理

核心逻辑转换 #

你原Pandas代码的逻辑是：按blockingId分组，每个组内筛选出country出现次数>1的行。用PySpark的窗口函数可以高效实现这个逻辑，无需循环遍历分组，完全适配百万级数据的处理场景。

完整代码实现 #

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

# 初始化SparkSession（未初始化时执行）
spark = SparkSession.builder.appName("CountryFilter").getOrCreate()

# 构造示例原始DataFrame（替换为你的实际数据）
old_df = spark.createDataFrame(
    [
        ("name_1", "block_1", "country_1"),
        ("name_2", "block_1", "country_2"),
        ("name_3", "block_2", "country_2"),
        ("name_4", "block_2", "country_2"),
    ],
    ["Name", "blockingId", "Country"]
)

# 定义窗口：按blockingId+Country的组合分组，统计每组出现次数
window_spec = Window.partitionBy("blockingId", "Country")

# 添加计数列、筛选目标行、清理临时列
new_df = old_df.withColumn("country_count", F.count("*").over(window_spec)) \
               .filter(F.col("country_count") > 1) \
               .drop("country_count")

# 查看结果
new_df.show()

代码解释 #

1. 窗口定义：Window.partitionBy("blockingId", "Country") 将数据按「blockingId分组+Country细分」切割，确保计数是每个blockingId组内单个Country的出现次数。
2. 计数列生成：F.count("*").over(window_spec) 为每一行添加临时列，值为当前行所属「blockingId+Country」组合的出现次数。
3. 筛选与清理：过滤出计数大于1的行，最后移除临时计数列，得到最终结果。

示例输出 #

运行后仅保留block_2下的两行数据：

+-------+---------+---------+
|   Name|blockingId|  Country|
+-------+---------+---------+
|name_3|  block_2|country_2|
|name_4|  block_2|country_2|
+-------+---------+---------+

扩展：分组内的更多操作 #

如果需要在blockingId分组内执行其他计算（比如组内统计、排序、排名等），可以灵活定义不同窗口：

• 仅按blockingId分组的窗口：Window.partitionBy("blockingId")，可用于计算每个blockingId组的总行数、均值等全局统计量。
• 嵌套组合逻辑：同时结合两种窗口实现多维度筛选，比如：

# 示例：先统计每个blockingId的总行数，再筛选Country出现次数>1的行
block_window = Window.partitionBy("blockingId")
new_df = old_df.withColumn("block_total", F.count("*").over(block_window)) \
               .withColumn("country_count", F.count("*").over(window_spec)) \
               .filter(F.col("country_count") > 1) \
               .drop("block_total", "country_count")

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