В Pyspark, как добавить список значений в качестве нового столбца в существующий Dataframe?

Question

У меня есть Pyspark Dataframe, подобный этому:

+--------+----+
|    col1|col2|
+--------+----+
|   Apple|   A|
|  Google|   G|
|Facebook|   F|
+--------+----+

У меня есть массив со значениями ["SFO", "LA", "NY C"]. Я хочу добавить этот массив в Dataframe как новый столбец, например:

#+--------+----+--------------+
#|    col1|col2|          col3|
#+--------+----+--------------+
#|   Apple|   A|SFO           |
#|  Google|   G|LA            |
#|Facebook|   F|NYC           |
#+--------+----+--------------+

Как это сделать в Pyspark? Я не должен использовать Pandas в своем решении.

Answer 1

Вы можете использовать функцию array и звездочку *, чтобы развернуть в ней свой список с помощью lit, чтобы поместить свой список в каждую строку новый столбец. Затем вы можете использовать вычисление row_number(), чтобы отправить результат этого на element_at. (Spark2.4+)

from pyspark.sql import functions as F
from pyspark.sql.window import Window

w=Window().orderBy("col3")
arr=["SFO","LA","NYC"]

df.withColumn("col3", F.array(*[F.lit(x) for x in arr]))\
  .withColumn("rownum", F.row_number().over(w))\
  .withColumn("col3", F.expr("""element_at(col3,rownum)""")).drop("rownum").show()

#+--------+----+----+
#|    col1|col2|col3|
#+--------+----+----+
#|   Apple|   A| SFO|
#|  Google|   G|  LA|
#|Facebook|   F| NYC|
#+--------+----+----+