使用 withColumn 向现有 PySpark DataFrame 添加两列

PySpark 中的 DataFrame 是一种分布式的数据集合，类似于关系型数据库中的表。在进行数据操作的时候，经常需要对 DataFrame 执行添加、删除、修改等操作，其中添加操作可以使用 withColumn 函数来完成。

withColumn 函数可以将 DataFrame 中的现有列加工处理后生成新的列，也可以直接添加新的列。以下我们将手把手教你如何使用 withColumn 向现有 PySpark DataFrame 添加两列。

1. 导入模块和创建 DataFrame

这里我们先导入必要的 Python 模块，再创建一个包含三列数据的 DataFrame，用于后文举例。

# 导入 PySpark 模块
from pyspark.sql.functions import col
from pyspark.sql import SparkSession

# 创建 SparkSession
spark = SparkSession.builder.appName("example").getOrCreate()

# 创建示例 DataFrame
data = [("A", "小明", 26), ("B", "小红", 25), ("C", "小刚", 27)]
df = spark.createDataFrame(data, ["id", "name", "age"])

# 显示 DataFrame
df.show()

运行以上代码，我们可以看到 DataFrame 的内容如下：

+---+----+---+
| id|name|age|
+---+----+---+
|  A|小明| 26|
|  B|小红| 25|
|  C|小刚| 27|
+---+----+---+

2. 使用 withColumn 添加列

假设我们现在需要向上面这个 DataFrame 添加两列（height 和 weight），分别存储每个人的身高和体重。那么我们可以使用 withColumn 函数来完成添加操作，如下：

# 使用 withColumn 添加 height 列，并填充为 180
df = df.withColumn("height", col("age") + 154)

# 使用 withColumn 添加 weight 列，并填充为 60.5
df = df.withColumn("weight", 60.5)

# 显示更改后的 DataFrame
df.show()

运行以上代码，我们可以看到 DataFrame 的内容如下：

+---+----+---+------+------+
| id|name|age|height|weight|
+---+----+---+------+------+
|  A|小明| 26|   180|  60.5|
|  B|小红| 25|   179|  60.5|
|  C|小刚| 27|   181|  60.5|
+---+----+---+------+------+

我们可以看到，使用 withColumn 函数向 DataFrame 中添加新的列非常简单，只需要使用 DataFrame 的 withColumn 函数，指定新列名和新列的取值表达式即可。在上面的例子中，我们使用 col 函数来引用 DataFrame 中的 age 列，并使用 + 运算符来进行加工处理生成新列 height。

3. 处理已经存在的列

在上文的例子中，我们展示了如何使用 withColumn 函数向 DataFrame 中添加新的列。如果我们要对已经存在的列进行加工处理呢？答案是可以的，我们也可以使用 withColumn 函数来实现。示例代码如下：

# 对 age 列做平均值加 1 的操作
df = df.withColumn("age", (col("age") + 1) * 1.5)

# 显示更改后的 DataFrame
df.show()

运行以上代码，我们可以看到 DataFrame 的内容如下：

+---+----+----+------+------+
| id|name| age|height|weight|
+---+----+----+------+------+
|  A|小明|  40|   180|  60.5|
|  B|小红|  39|   179|  60.5|
|  C|小刚|  42|   181|  60.5|
+---+----+----+------+------+

我们可以看到，使用 withColumn 函数对已经存在的列进行加工处理非常方便，只需要在 withColumn 函数中指定已经存在的列名和新列的取值表达式即可。

总之，使用 withColumn 函数进行 DataFrame 操作非常方便，可谓是 PySpark 中的一大利器。了解了以上示例后，您应该可以轻松地使用 withColumn 函数来完成常见的 DataFrame 操作。