как я могу распаковать столбец списка типов в pyspark

Question

У меня есть dataframe в pyspark, у df есть столбец типа string array, поэтому мне нужно сгенерировать новый столбец с заголовком списка, а также мне нужен другой столбец с concat из хвостового списка.

это мой оригинальный фрейм данных:

pyspark> df.show()
+---+------------+
| id|     lst_col|
+---+------------+
|  1|[a, b, c, d]|
+---+------------+


pyspark> df.printSchema()
root
 |-- id: integer (nullable = false)
 |-- lst_col: array (nullable = true)
 |    |-- element: string (containsNull = true)

, и мне нужно сгенерировать что-то вроде этого:

pyspark> df2.show()
+---+--------+---------------+
| id|lst_head|lst_concat_tail|
+---+--------+---------------+
|  1|       a|          b,c,d|
+---+--------+---------------+

Answer 1

Для Spark 2.4+ вы можете использовать функции element_at, slice и size для массивов:

df.select("id",
          element_at("lst_col", 1).alias("lst_head"),
          expr("slice(lst_col, 2, size(lst_col))").alias("lst_concat_tail")
         ).show()

Дает:

+---+--------+---------------+
| id|lst_head|lst_concat_tail|
+---+--------+---------------+
|  1|       a|      [b, c, d]|
+---+--------+---------------+