Извлечение массива NumPy из Pyspark Dataframe

Question

У меня есть датафрейм gi_man_df, где группа может быть:

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

и я ожидаю numy nd_array, т.е.gi_man_array:

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

где значения rand_double после применения пивота.

Я попробовал следующие 2 подхода:
Первый: я поворачиваю 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 подходов верным способом достижения желаемого результата, если да, то как я могу продолжить, чтобы получить желаемый результат?
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()