基于Pandas/PySpark实现列重排并按重复列标记状态颜色

先构造示例数据集方便演示：

import pandas as pd

data = {
    'A': ['X1', 'X1', 'X2', 'X3', 'X2', 'X4', 'X1'],
    'B': ['b1', 'b1', 'b2', 'b3', 'b2', 'b4', 'b5'],
    'C': ['c1', 'c1', 'c2', 'c3', 'c2', 'c4', 'c5'],
    'D': ['d1', 'd1', 'd2', 'd3', 'd2', 'd4', 'd5'],
    'status_color': ['red', 'blue', 'green', 'yellow', 'orange', 'purple', 'gray']
}
df = pd.DataFrame(data)

实现步骤 #

1. 标记B、C、D组合的重复行：用duplicated参数keep=False，让所有重复组的行都被标记为True
2. 遍历A列的唯一值，为每个值生成新列：当该行的A等于目标值且B/C/D重复时，填入status_color，否则填0

代码实现 #

# 标记B/C/D重复的行
df['is_duplicated'] = df.duplicated(subset=['B', 'C', 'D'], keep=False)

# 生成每个A值对应的新列
for a_val in df['A'].unique():
    df[a_val] = df.apply(
        lambda row: row['status_color'] if row['A'] == a_val and row['is_duplicated'] else 0,
        axis=1
    )

# 可选：删除临时标记列
df = df.drop('is_duplicated', axis=1)

先构造示例Spark DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, count
from pyspark.sql.window import Window

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

data = [
    ('X1', 'b1', 'c1', 'd1', 'red'),
    ('X1', 'b1', 'c1', 'd1', 'blue'),
    ('X2', 'b2', 'c2', 'd2', 'green'),
    ('X3', 'b3', 'c3', 'd3', 'yellow'),
    ('X2', 'b2', 'c2', 'd2', 'orange'),
    ('X4', 'b4', 'c4', 'd4', 'purple'),
    ('X1', 'b5', 'c5', 'd5', 'gray')
]
df = spark.createDataFrame(data, ['A', 'B', 'C', 'D', 'status_color'])

实现步骤 #

1. 用窗口函数计算每个B/C/D组合的行数，判断是否重复（行数>1）
2. 遍历A的唯一值，用when条件生成对应新列：当A等于目标值且组合重复时，取status_color，否则填0

代码实现 #

# 计算每个B/C/D组合的行数，标记是否重复
window = Window.partitionBy('B', 'C', 'D')
df = df.withColumn('group_count', count('*').over(window))
df = df.withColumn('is_duplicated', col('group_count') > 1)

# 获取A的唯一值集合
a_values = [row['A'] for row in df.select('A').distinct().collect()]

# 生成每个A对应的新列
for a_val in a_values:
    df = df.withColumn(
        a_val,
        when((col('A') == a_val) & col('is_duplicated'), col('status_color')).otherwise(0)
    )

# 可选：删除临时计算列
df = df.drop('group_count', 'is_duplicated')

# 查看结果
df.show()

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