在 Spark 数据框中添加列 - Python

在 Spark 数据框中添加列是一项常见的任务。本文将介绍如何使用 Python 在 Spark 数据框中添加新列。

创建 Spark 数据框

在开始添加新列之前，我们需要创建一个 Spark 数据框。我们可以使用以下代码创建一个简单的数据框：

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("Add Column").getOrCreate()

data = [("Alice", 25), ("Bob", 30), ("Charlie", 35)]
columns = ["Name", "Age"]

df = spark.createDataFrame(data, columns)

这会创建一个如下所示的数据框：

+-------+---+
|   Name|Age|
+-------+---+
|  Alice| 25|
|    Bob| 30|
|Charlie| 35|
+-------+---+

添加新列

我们使用 withColumn 方法来添加新列。此方法需要两个参数：列名和要添加的列的值。我们可以使用以下代码来添加一个新的工资列：

from pyspark.sql.functions import col

df = df.withColumn("Salary", col("Age") * 1000)

这将在我们的数据框中添加一个名为 "Salary" 的新列，并将每个人的工资设置为他们的年龄乘以 1000。我们可以使用以下代码检查新列是否添加成功：

df.show()

输出：

+-------+---+------+
|   Name|Age|Salary|
+-------+---+------+
|  Alice| 25| 25000|
|    Bob| 30| 30000|
|Charlie| 35| 35000|
+-------+---+------+

我们还可以使用多个列执行更复杂的计算。例如，我们可以使用以下代码添加一个名为 "Bonus" 的新列，该列为每个人的工资加上 5000：

df = df.withColumn("Bonus", col("Salary") + 5000)

此时，我们的数据框如下所示：

+-------+---+------+-----+
|   Name|Age|Salary|Bonus|
+-------+---+------+-----+
|  Alice| 25| 25000|30000|
|    Bob| 30| 30000|35000|
|Charlie| 35| 35000|40000|
+-------+---+------+-----+

结论

在 Spark 数据框中添加列非常简单。我们只需要使用 withColumn 方法并提供新列的名称和值即可。在执行计算时，我们可以使用现有的列来创建新列。