Добавьте разные массивы из numpy в каждый ряд данных - PullRequest
Новая
       35

Добавьте разные массивы из numpy в каждый ряд данных

0 голосов
/ 04 октября 2019

У меня есть датафрейм SparkSQL и двумерная матрица. У них одинаковое количество рядов. Я намерен добавить каждый отдельный массив из простой матрицы в качестве нового столбца в существующий фрейм данных 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]

Идентификатор столбца соответствует порядку записей в матричной матрице.

Интересно, есть ли эффективный способ реализовать это?

1 Ответ

2 голосов
/ 04 октября 2019

Создайте DataFrame из вашей матрицы numpy и добавьте столбец Id, чтобы указать номер строки. Затем вы можете присоединиться к исходному фрейму данных PySpark в столбце Id. 

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