将不同的数组从 numpy 添加到数据帧的每一行

Question

我有一个 SparkSQL 数据框和 2D numpy 矩阵。它们具有相同的行数。我打算将 numpy 矩阵中的每个不同数组作为新列添加到现有的 PySpark 数据框中。这样每一行添加的列表就不同了。

例如，PySpark 数据框是这样的

| Id     | Name   |
| ------ | ------ |
| 1      | Bob    |
| 2      | Alice  |
| 3      | Mike   |

而numpy矩阵是这样的

[[2, 3, 5]
 [5, 2, 6]
 [1, 4, 7]]

生成的预期数据框应该是这样的

| Id     | Name   | customized_list
| ------ | ------ | ---------------
| 1      | Bob    |   [2, 3, 5]
| 2      | Alice  |   [5, 2, 6]
| 3      | Mike   |   [1, 4, 7]

Id 列对应于 numpy 矩阵中条目的顺序。

我想知道有没有有效的方法来实现这个？

Answer 1

从您的numpy 矩阵创建一个DataFrame，并添加一个Id 列来指示行号。然后，您可以在 Id 列上加入原始 PySpark DataFrame。

import numpy as np
a = np.array([[2, 3, 5], [5, 2, 6], [1, 4, 7]])
list_df = spark.createDataFrame(enumerate(a.tolist(), start=1), ["Id", "customized_list"])
list_df.show()
#+---+---------------+
#| Id|customized_list|
#+---+---------------+
#|  1|      [2, 3, 5]|
#|  2|      [5, 2, 6]|
#|  3|      [1, 4, 7]|
#+---+---------------+

这里我使用enumerate(..., start=1)添加行号。

现在只需进行内部连接：

df.join(list_df, on="Id", how="inner").show()
#+---+-----+---------------+
#| Id| Name|customized_list|
#+---+-----+---------------+
#|  1|  Bob|      [2, 3, 5]|
#|  3| Mike|      [1, 4, 7]|
#|  2|Alice|      [5, 2, 6]|
#+---+-----+---------------+