Как разделить столбец со значениями, разделенными запятыми, и сохранить их в массиве в Dataframe PySpark? Как дано ниже

Question

У меня есть фрейм данных PySpark со столбцом, который содержит значения, разделенные запятыми. Количество значений, содержащихся в столбце, является фиксированным (скажем, 4). Пример:

+------------------------+
|col1                    |
+------------------------+
|1,val1, val4            |
|2,val1                  |
|3,val1, val2, val3      |
|4,val1, val2, val3, val4|
+------------------------+

Теперь я хочу, чтобы он был разбит на 2 столбца, как показано ниже

+----+------------------------+
|col1|col2                    |
+----+------------------------+
|   1|[val1, val4]            |
|   2|[val1]                  |
|   3|[val1, val2, val3]      |
|   4|[val1, val2, val3, val4]|
+----+------------------------+

Как это можно сделать?

Answer 1

Вы можете реализовать это, используя slice и split :

from pyspark.sql.functions import col, split, slice

array_len = 4
df.withColumn("ar", split(col("col1"), ",")) \
  .select(
     col("ar")[0].alias("col1"), 
     slice(col("ar"), 2, array_len).alias("col2")
  )

# +----+---------------------------+
# |col1|col2                       |
# +----+---------------------------+
# |1   |[val1,  val4]              |
# |2   |[val1]                     |
# |3   |[val1,  val2,  val3]       |
# |4   |[val1,  val2,  val3,  val4]|
# +----+---------------------------+

Сначала мы разбиваем и сохраняем массив в ar, затем мы используем select для получить первый элемент массива с помощью col("ar")[0], а остальную часть массива - с помощью slice(col("ar"), 2, array_len), который вернет все элементы, кроме первого.