在 PySpark 数据框中添加列总和作为新列

PySpark 是 Apache Spark 用 Python API 编程的方式，它提供了一个容易使用且效率高的分布式数据处理框架。在 PySpark 中处理数据是非常常见的操作，本文将介绍如何在 PySpark 数据框中添加列总和作为新列。

准备数据

首先，我们需要创建一个 PySpark 数据框并添加一些数据：

# 导入必要的库
from pyspark.sql import SparkSession
from pyspark.sql.functions import col

# 创建 SparkSession
spark = SparkSession.builder.appName("Add Column Total Sum").getOrCreate()

# 创建数据框
data = [("Alice", 25, 1000), ("Bob", 30, 1500), ("Charlie", 35, 2000)]
columns = ["name", "age", "salary"]
df = spark.createDataFrame(data, columns)

# 显示数据框
df.show()

输出结果：

+-------+---+------+
|   name|age|salary|
+-------+---+------+
|  Alice| 25|  1000|
|    Bob| 30|  1500|
|Charlie| 35|  2000|
+-------+---+------+

这里我们创建了一个有三列的数据框，其中包含每个人的名字、年龄和薪水。

添加列总和

接下来，我们将使用 withColumn 函数添加一个新列 total，用于记录每个人年龄和薪水的总和：

# 添加新列
df_with_total = df.withColumn("total", col("age") + col("salary"))

# 显示数据框
df_with_total.show()

输出结果：

+-------+---+------+-----+
|   name|age|salary|total|
+-------+---+------+-----+
|  Alice| 25|  1000| 1025|
|    Bob| 30|  1500| 1530|
|Charlie| 35|  2000| 2035|
+-------+---+------+-----+

我们使用 withColumn 函数将 total 列添加到数据框中。这个函数接受两个参数：第一个参数是新列的名称，第二个参数是一个表达式，用于计算新列的值。在这个表达式中，我们使用了 col 函数来获取 age 和 salary 列的值，并将它们相加得到总和。

结论

这样，我们就成功地在 PySpark 数据框中添加了列总和作为新列。PySpark 提供了许多方便的函数来处理数据，使得数据处理变得更加容易和高效。在使用 PySpark 进行数据处理时，请注意其特殊的语法和函数，以及分布式计算的局限性。