在PySpark或Pandas中实现DataFrame列插值并按指定列降序排序

我来帮你搞定这个需求！不管用Pandas还是PySpark，都能实现把my_column的值转成独立列，同时按每个some_id对应的int_column值降序排序的效果，下面分别给出两种工具的实现步骤：

Pandas 实现 #

首先我们先构造示例数据，然后按步骤操作：

步骤1：构造示例DataFrame #

import pandas as pd

data = {
    'some_id': ['xx1', 'xx1', 'xx1', 'xx2', 'xx2', 'xx2', 'xx3', 'xx3', 'xx3'],
    'my_column': ['id_1', 'id_2', 'id_3', 'id_1', 'id_2', 'id_3', 'id_1', 'id_2', 'id_3'],
    'int_column': [3, 4, 5, 6, 1, 3, 4, 2, 7]
}
df = pd.DataFrame(data)

步骤2：按组内int_column降序排序 #

先对每个some_id分组，组内按照int_column从大到小排序：

# 先按some_id升序（保持组的顺序），再按int_column降序排序
sorted_df = df.sort_values(by=['some_id', 'int_column'], ascending=[True, False])

步骤3：行转列（将my_column转为独立列） #

用pivot方法把my_column的每个值转成单独的列，int_column作为对应列的值：

pivoted_df = sorted_df.pivot(index='some_id', columns='my_column', values='int_column')

如果需要调整列的顺序为全局int_column的降序，可以额外添加一步：

# 计算每个my_column对应的int_column平均值，按平均值降序排列列
column_order = pivoted_df.mean().sort_values(ascending=False).index
pivoted_df = pivoted_df[column_order]

PySpark 实现 #

如果你用的是PySpark，步骤类似，不过需要用Spark的API来操作：

步骤1：初始化SparkSession并构造示例DataFrame #

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

# 初始化SparkSession
spark = SparkSession.builder.appName("pivot_sort").getOrCreate()

# 构造示例数据
data = [
    ("xx1", "id_1", 3),
    ("xx1", "id_2", 4),
    ("xx1", "id_3", 5),
    ("xx2", "id_1", 6),
    ("xx2", "id_2", 1),
    ("xx2", "id_3", 3),
    ("xx3", "id_1", 4),
    ("xx3", "id_2", 2),
    ("xx3", "id_3", 7)
]
df = spark.createDataFrame(data, ["some_id", "my_column", "int_column"])

步骤2：按组内int_column降序排序 #

用窗口函数给每个some_id组内的行按int_column降序排名，然后按排名排序：

# 定义窗口：按some_id分组，按int_column降序排序
window = Window.partitionBy("some_id").orderBy(F.col("int_column").desc())

# 添加排名列，然后按排名排序
sorted_df = df.withColumn("rank", F.row_number().over(window)).orderBy("rank")

步骤3：行转列（pivot操作） #

用pivot方法将my_column转为独立列，同时聚合int_column的值：

pivoted_df = sorted_df.groupBy("some_id").pivot("my_column").agg(F.first("int_column"))

步骤4：调整列顺序（可选） #

如果需要让列按全局int_column的降序排列，可以计算每个列的平均值，然后重新排列：

# 计算每个my_column对应的int_column平均值
col_avg = pivoted_df.groupBy().avg().collect()[0]
# 生成按平均值降序的列顺序
column_order = sorted([col[4:] for col in col_avg.asDict().keys()], 
                      key=lambda x: col_avg[f"avg({x})"], reverse=True)
# 重新排列列
pivoted_df = pivoted_df.select("some_id", *column_order)

最后展示结果：

pivoted_df.show()

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