从 Pyspark Dataframe 中提取 numpy 数组

Question

我有一个数据框gi_man_df，其中组可以是n：

+------------------+-----------------+--------+--------------+
|           group  |           number|rand_int|   rand_double|
+------------------+-----------------+--------+--------------+
|          'GI_MAN'|                7|       3|         124.2|
|          'GI_MAN'|                7|      10|        121.15|
|          'GI_MAN'|                7|      11|         129.0|
|          'GI_MAN'|                7|      12|         125.0|
|          'GI_MAN'|                7|      13|         125.0|
|          'GI_MAN'|                7|      21|         127.0|
|          'GI_MAN'|                7|      22|         126.0|
+------------------+-----------------+--------+--------------+

我期待一个 numpy nd_array 即 gi_man_array:

[[[124.2],[121.15],[129.0],[125.0],[125.0],[127.0],[126.0]]]

其中应用主元后的 rand_double 值。

我尝试了以下两种方法：
首先： 我按如下方式旋转 gi_man_df：

gi_man_pivot = gi_man_df.groupBy("number").pivot('rand_int').sum("rand_double")

我得到的输出是：

Row(number=7, group=u'GI_MAN', 3=124.2, 10=121.15, 11=129.0, 12=125.0, 13=125.0, 21=127.0, 23=126.0)

但这里的问题是获得所需的输出，我无法将其转换为矩阵然后再次转换为 numpy 数组。

第二： 我使用以下方法在数据帧本身中创建了向量：

assembler = VectorAssembler(inputCols=["rand_double"],outputCol="rand_double_vector")

gi_man_vector = assembler.transform(gi_man_df)
gi_man_vector.show(7)

我得到了以下输出：

+----------------+-----------------+--------+--------------+--------------+
|           group|           number|rand_int|   rand_double| rand_dbl_Vect|
+----------------+-----------------+--------+--------------+--------------+
|          GI_MAN|                7|       3|         124.2|       [124.2]|
|          GI_MAN|                7|      10|        121.15|      [121.15]|
|          GI_MAN|                7|      11|         129.0|       [129.0]|
|          GI_MAN|                7|      12|         125.0|       [125.0]|
|          GI_MAN|                7|      13|         125.0|       [125.0]|
|          GI_MAN|                7|      21|         127.0|       [127.0]|
|          GI_MAN|                7|      22|         126.0|       [126.0]|
+----------------+-----------------+--------+--------------+--------------+

但问题是我无法将其旋转到 rand_dbl_Vect。

所以我的问题是：
1.这两种方法中的任何一种都是实现所需输出的正确方法吗？如果是，那么我该如何进一步进行以获得所需的结果？
2. 我还可以采取什么其他方法来使代码最优且性能良好？

Answer 1

这个

import numpy as np
np.array(gi_man_df.select('rand_double').collect())

产生

array([[ 124.2 ],
       [ 121.15],
       .........])

Answer 2

要将 Spark df 转换为 numpy 数组，请首先将其转换为 pandas，然后应用 to_numpy() 函数。

spark_df.select(<list of columns needed>).toPandas().to_numpy()

Answer 3

这个基于@data_steve答案的解决方案可以提高内存效率，但需要更长的时间：

import numpy as np
np.fromiter( gi_man_df.select('rand_double').toLocalIterator(), dtype=float )

因为它不会首先创建本地数据帧，然后创建另一个 numpy 数组，而是逐一读取值来构建数组。我观察了 RAM 消耗，这似乎正是正在发生（未发生）的情况。

也许您可以指定更合适的

type

。