如何将numpy二维数组存入PySpark DataFrame单个单元格

方案实现 #

场景1：将完整二维数组存入单个行单元格 #

如果你需要把整个3x3的numpy二维数组作为唯一一行的单个字段值，直接将numpy数组转为嵌套列表后创建DataFrame即可：

import pyspark.sql.functions as f
import numpy as np
from pyspark.sql.types import ArrayType, DoubleType

# 定义原始numpy二维数组
np_2d_arr = np.array([
    [0.        , 0.67235401, 0.35767577],
    [0.67235401, 0.        , 0.2981656 ],
    [0.35767577, 0.2981656 , 0.        ]
])

# 嵌套列表传入创建DataFrame，指定嵌套数组类型
df = spark.createDataFrame(
    [(np_2d_arr.tolist(),)], 
    schema=ArrayType(ArrayType(DoubleType()))
).withColumnRenamed("_1", "array")

df.show(truncate=False)

场景2：二维数组每行对应DataFrame一行，单字段存该行一维数组 #

如果你需要原始二维数组的每一行对应PySpark DataFrame的一行，每行的array字段存储该行的数值数组：

import pyspark.sql.functions as f
import numpy as np

df = spark.createDataFrame(np.array([
    [0.        , 0.67235401, 0.35767577],
    [0.67235401, 0.        , 0.2981656 ],
    [0.35767577, 0.2981656 , 0.        ]
]))

# 用array函数合并所有列得到行数组
df = df.withColumn('array', f.array(*df.columns))

df.show(truncate=False)

原代码问题说明 #

• 写完df = spark.createDataFrame(...)后调用了不存在的变量sms，会直接抛出变量未定义错误
• 你编写的TRANSFORM表达式逻辑不符合需求，不需要用高阶函数即可实现对应效果，多余的写法反而会导致逻辑混乱。

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