PySpark 使用字典创建新列

在本文中，我们将介绍如何使用 PySpark 中的字典创建新列。在数据处理过程中，有时候我们需要根据一个字典表进行映射，将某一列的值映射为另一列的值。PySpark 提供了方便的函数和方法来实现这个功能。

准备工作

在使用 PySpark 创建新列之前，我们需要导入相应的库，并创建一个 SparkSession 对象。可以通过以下代码来完成准备工作：

from pyspark.sql import SparkSession

# 创建 SparkSession 对象
spark = SparkSession.builder.appName("CreateNewColumn").getOrCreate()

创建数据框

在示例中，我们将以一个示例数据框为例进行操作。假设我们有一个包含用户 ID 和对应姓名的数据框，格式如下所示：

ID	Name
1	Alice
2	Bob
3	Charlie
4	David

我们可以使用以下代码创建这个数据框：

from pyspark.sql.types import StructType, StructField, IntegerType, StringType

# 定义模式
schema = StructType([
    StructField("ID", IntegerType(), True),
    StructField("Name", StringType(), True)
])

# 创建数据
data = [
    (1, "Alice"),
    (2, "Bob"),
    (3, "Charlie"),
    (4, "David")
]

# 创建数据框
df = spark.createDataFrame(data, schema)

现在我们已经创建了一个名为 df 的数据框。

使用字典创建新列

假设我们有一个字典表，包含用户 ID 和对应的年龄信息。我们想要根据字典表将数据框 df 中的用户 ID 映射为对应的年龄，并创建一个新列 Age。我们可以使用 PySpark 中的 withColumn 方法和 udf 函数来实现这个功能。

首先，我们需要将字典转换为 PySpark 中的字典类型。然后，我们可以创建一个用户自定义函数（UDF），将字典中对应的年龄取出并返回。接下来，我们可以使用 withColumn 方法来创建新列，并将 UDF 应用在其中。

以下是具体的代码示例：

from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType

# 定义字典
age_dict = {1: 25, 2: 30, 3: 35, 4: 40}

# 将字典转换为 PySpark 中的字典类型
age_dict_udf = udf(lambda id: age_dict.get(id), IntegerType())

# 创建新列
df_with_age = df.withColumn("Age", age_dict_udf(df["ID"]))

# 查看结果
df_with_age.show()

运行以上代码，我们将获得一个新的数据框 df_with_age，其中包含了根据字典映射后的年龄信息：