如何在Pyspark DataFrame中按指定规则计算对角线值总和

逻辑说明 #

• w1对应对角线：取行号与列序号匹配的单元格值求和，即第1行取Col1、第2行取Col2、……、第5行取Col5
• w2对应对角线：从第3行第2列（你提到的B3位置）开始取对角线值，即第3行取Col2、第4行取Col3、第5行取Col4、第6行取Col5

完整代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import monotonically_increasing_id, row_number, col, sum, when
from pyspark.sql.window import Window

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

# 构造测试数据
test_data = [
    ("w1", 161.8, 161.8, 161.8, 161.8, 161.8),
    ("w2", 112.0, 112.0, 112.0, 112.0, 112.0),
    ("w3", 157.7, 157.7, 157.7, 157.7, 157.7),
    ("w4", 146.1, 146.1, 146.1, 146.1, 146.1),
    ("w5", 193.4, 193.4, 193.4, 193.4, 193.4),
    ("w6", 44.7, 44.7, 44.7, 44.7, 44.7),
    ("w7", 105.3, 105.3, 105.3, 105.3, 105.3)
]
df = spark.createDataFrame(test_data, schema=["w", "Col1", "Col2", "Col3", "Col4", "Col5"])

# 生成连续行号，保证和原始数据行顺序一致
row_window = Window.orderBy(monotonically_increasing_id())
df_with_rowid = df.withColumn("row_id", row_number().over(row_window))

# 提取两个对角线的对应单元格值
diag_value_df = df_with_rowid.withColumn(
    "w1_diag_value",
    when(col("row_id") == 1, col("Col1"))
    .when(col("row_id") == 2, col("Col2"))
    .when(col("row_id") == 3, col("Col3"))
    .when(col("row_id") == 4, col("Col4"))
    .when(col("row_id") == 5, col("Col5"))
).withColumn(
    "w2_diag_value",
    when(col("row_id") == 3, col("Col2"))
    .when(col("row_id") == 4, col("Col3"))
    .when(col("row_id") == 5, col("Col4"))
    .when(col("row_id") == 6, col("Col5"))
)

# 计算求和结果
sum_result = diag_value_df.select(
    sum("w1_diag_value").alias("w1对角线总和"),
    sum("w2_diag_value").alias("w2对角线总和")
)

# 输出结果
sum_result.show()

输出示例 #

+-------------+-------------+
|w1对角线总和|w2对角线总和|
+-------------+-------------+
|        771.0|        541.9|
+-------------+-------------+

如果后续对角线的起始位置、列数有调整，只需要修改when的匹配规则即可，不需要调整整体逻辑。

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