PySpark数据清洗：识别重复列并生成指定结果列

现有如下PySpark DataFrame：

+----+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+
| ID | Column_ID_1 | Column_txt_1 | Column_ID_2 | Column_txt_2 | Column_ID_3 | Column_txt_3 | Column_ID_4 | Column_txt_4 | Column_ID_5 | Column_txt_5 | Column_ID_6 | Column_txt_6 | Column_ID_7 | Column_txt_7 |
+----+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+
|  1 |         100 | abc          |       10025 | lkegls       |       64651 | wlejgnlwnr   |       65132 | wejgnn       |       65132 | wejgnn       |       65132 | wejgnn       |       65132 | wejgnn       |
|  2 |        1001 | def          |      648613 | esglknsl     |     3213213 | weklgnlkwn   |       51741 | oiqjerml     |  6565468165 | welkmlwe     |  6565468165 | welkmlwe     |  6565468165 | welkmlwe     |
|  3 |        1002 | hijlkefn     |        5654 | lengl        |      513220 | wlegnlw      |      513220 | wlegnlw      |      513220 | wlegnlw      |      513220 | wlegnlw      |      513220 | wlegnlw      |
+----+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+-------------+--------------+

需要识别每个ID对应的首次出现重复的Column_ID列名及其对应的Column_txt值，生成如下结果：

+----+------------------+----------------------+
| ID | Repeating_Column | Repeating Column_txt |
+----+------------------+----------------------+
|  1 | Column_ID_4      | wejgnn               |
|  2 | Column_ID_5      | welkmlwe             |
|  3 | Column_ID_3      | wlegnlw              |
+----+------------------+----------------------+

步骤1：构建测试DataFrame #

先还原示例中的测试数据：

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

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

data = [
    (1, 100, "abc", 10025, "lkegls", 64651, "wlejgnlwnr", 65132, "wejgnn", 65132, "wejgnn", 65132, "wejgnn", 65132, "wejgnn"),
    (2, 1001, "def", 648613, "esglknsl", 3213213, "weklgnlkwn", 51741, "oiqjerml", 6565468165, "welkmlwe", 6565468165, "welkmlwe", 6565468165, "welkmlwe"),
    (3, 1002, "hijlkefn", 5654, "lengl", 513220, "wlegnlw", 513220, "wlegnlw", 513220, "wlegnlw", 513220, "wlegnlw", 513220, "wlegnlw")
]

columns = [
    "ID", "Column_ID_1", "Column_txt_1", "Column_ID_2", "Column_txt_2",
    "Column_ID_3", "Column_txt_3", "Column_ID_4", "Column_txt_4",
    "Column_ID_5", "Column_txt_5", "Column_ID_6", "Column_txt_6",
    "Column_ID_7", "Column_txt_7"
]

df = spark.createDataFrame(data, columns)

步骤2：UDF实现查找首次重复项 #

核心思路是将每对Column_ID_n和Column_txt_n打包成结构体数组，遍历数组找出第一个与已出现过的键重复的项：

# 生成包含(ID值、txt值、列名)的结构体数组
pair_cols = [
    F.struct(
        F.col(f"Column_ID_{i}").alias("id_val"),
        F.col(f"Column_txt_{i}").alias("txt_val"),
        F.lit(f"Column_ID_{i}").alias("col_name")
    ) for i in range(1, 8)
]
df_with_pairs = df.withColumn("pairs", F.array(*pair_cols))

# 定义UDF找出第一个重复项
def find_first_repeat(pairs):
    seen_keys = set()
    for pair in pairs:
        current_key = (pair.id_val, pair.txt_val)
        if current_key in seen_keys:
            return (pair.col_name, pair.txt_val)
        seen_keys.add(current_key)
    return (None, None)

find_first_repeat_udf = F.udf(find_first_repeat, "struct<Repeating_Column:string, `Repeating Column_txt`:string>")

# 应用UDF并提取结果
result_df = df_with_pairs.withColumn("repeat_info", find_first_repeat_udf(F.col("pairs"))) \
    .select("ID", "repeat_info.*")

# 展示结果
result_df.show()

步骤3：纯PySpark函数实现（无UDF） #

如果偏好不用UDF，可通过数组展开+窗口函数实现：

# 生成带索引的(ID值、txt值、列名)结构体数组
pair_cols_with_idx = [
    F.struct(
        F.lit(i).alias("idx"),
        F.col(f"Column_ID_{i}").alias("id_val"),
        F.col(f"Column_txt_{i}").alias("txt_val"),
        F.lit(f"Column_ID_{i}").alias("col_name")
    ) for i in range(1, 8)
]
df_with_pairs = df.withColumn("pairs", F.array(*pair_cols_with_idx))

# 展开数组并标记重复项
df_exploded = df_with_pairs.select("ID", F.explode("pairs").alias("pair")) \
    .withColumn("count", F.count("*").over(F.window.partitionBy("ID", "pair.id_val", "pair.txt_val"))) \
    .filter(F.col("count") > 1)

# 找出每个ID最小重复索引对应的项
result_df = df_exploded.groupBy("ID") \
    .agg(F.min("pair.idx").alias("min_idx")) \
    .join(df_exploded, on=["ID", "min_idx"], how="left") \
    .select("ID", "pair.col_name", "pair.txt_val") \
    .withColumnRenamed("col_name", "Repeating_Column") \
    .withColumnRenamed("txt_val", "Repeating Column_txt")

result_df.show()

验证结果 #

两种方法执行后，输出均符合需求：

+---+------------------+----------------------+
| ID|Repeating_Column  |Repeating Column_txt   |
+---+------------------+----------------------+
|  1|Column_ID_4       |wejgnn                |
|  2|Column_ID_5       |welkmlwe              |
|  3|Column_ID_3       |wlegnlw               |
+---+------------------+----------------------+

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