PySpark – 从两列数据创建字典

在PySpark中，有时需要从两列数据创建字典，可以使用map和collectAsMap方法。本文将介绍如何使用这两个方法来创建字典。

1. 从两个列创建字典

假设我们有一个名为data的数据框，其中包含两个列：key和value。我们要将这两个列转换为字典。下面是一个示例。

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.master('local').appName('Dictionary from two columns').getOrCreate()

# 创建示例数据集
data = [(1, 'one'), (2, 'two'), (3, 'three')]
df = spark.createDataFrame(data, ['key', 'value'])

# 从两个列创建字典
dict_data = df.rdd.map(lambda x: (x[0], x[1])).collectAsMap()

print(dict_data)

输出如下:

{1: 'one', 2: 'two', 3: 'three'}

首先我们创建了一个SparkSession对象，然后使用createDataFrame方法从一个列表data中创建了一个名为df的数据框。接下来，我们使用rdd属性将数据框转换为RDD，然后使用map方法将每个元组转换为一个键值对，并使用collectAsMap方法将这些键值对收集到一个字典中。

2. 将两个列合并到一个字典中

还可以将两个列合并到一个字典中，可以使用create_map函数来实现。下面是一个示例。

from pyspark.sql import functions as F

# 创建示例数据集
data = [(1, 'one'), (2, 'two'), (3, 'three')]
df = spark.createDataFrame(data, ['key', 'value'])

# 将两个列合并到一个字典中
df = df.withColumn('dict', F.create_map(df['key'], df['value']))

df.show(truncate=False)

输出如下:

+---+-----+----------------+
|key|value|dict            |
+---+-----+----------------+
|1  |one  |[1 -> one]      |
|2  |two  |[2 -> two]      |
|3  |three|[3 -> three]    |
+---+-----+----------------+

首先，我们创建了一个名为df的数据框,然后使用withColumn方法将两个列合并为一个dict列，并使用create_map函数将这两列转换为一个字典。