PySpark中如何将数组列转换为含静态列表索引的结构体

要实现将数组列与静态键列表映射为可通过键查询的映射列，可利用Spark的create_map函数生成键值对映射，具体步骤如下：

代码实现 #

1. 创建示例DataFrame并定义静态键列表：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, create_map, lit

spark = SparkSession.builder.appName("array_to_map").getOrCreate()

# 示例输入DataFrame
data = [(["a","b","c"],), (["d","e","f"],)]
df = spark.createDataFrame(data, ["col1"])
keys = [1, 1.5, 2]

2. 生成create_map所需的交替键值对参数：
   create_map要求参数按「键、值、键、值」的顺序传入，遍历静态键列表，将每个键转为Spark字面量，再对应数组列的对应位置元素：

map_args = []
for idx, key in enumerate(keys):
    map_args.append(lit(key))  # 将静态键转为Spark可识别的字面量
    map_args.append(col("col1").getItem(idx))  # 取数组对应位置的元素作为值

3. 创建映射列并验证查询：

# 添加new_column列，类型为Spark MapType
df_result = df.withColumn("new_column", create_map(*map_args))

# 验证通过键查询值的操作（两种方式均支持）
# 方式1：使用getItem方法
df_result.select(col("new_column").getItem(1.5)).show()
# 方式2：使用[]语法（Spark 3.x版本支持）
df_result.select(col("new_column")[1]).show()

关键说明 #

• 为何不用struct？Spark的Struct类型字段名必须是字符串，而你的键包含数字/浮点数，无法直接作为struct字段名，且struct不支持动态键查询。
• 关于键类型：keys包含整数和浮点数时，Spark会自动将键类型统一为DoubleType，查询整数1实际匹配的是1.0；若需要严格类型匹配，建议将keys统一为浮点数（如[1.0, 1.5, 2.0]）。
• 该方案完全适配Spark 3.3.0版本，生成的new_column是标准Map类型，支持任意合法键的查询操作。

内容的提问来源于stack exchange，提问作者Maya