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]|
#+---+-----+---------------+

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

1 个答案: