Повторите элементы массива k раз в pyspark

Question

Пример данных DF со столбцом массива:

+--------------------+
|         COL1       |
+--------------------+
|[A, B, C]           |
+--------------------+

Я пытался использовать array_repeat , но мой тип данных уже имеет ArrayType

df.select(array_repeat(df.COL1, 2).alias('C')).collect()

Вывод:

+----------------------+
|         C            |
+----------------------+
|[[A, B, C],[A, B, C]] |
+----------------------+

Создан массив массивов

 C: array (nullable = false)
 |    |-- element: array (containsNull = true)
 |    |    |-- element: string (containsNull = false)

Ожидаемый результат:

+----------------------+
|         C            |
+----------------------+
|[A, B, C, A, B, C]    |
+----------------------+

Answer 1

В Spark 2.4.0 появилась функция flatten(), которая удаляет один уровень вложенности массива. Учитывая, что 2.4.0 - это та же версия Spark, которая представила array_repeat(), вам просто нужно заключить вызов в array_repeat() в flatten():

>>> df.select(flatten(array_repeat(df.COL1, 2)).alias("C")).show(1, False)
+------------------------+
|C                       |
+------------------------+
|[A, B, C, A, A, B, C, A]|
+------------------------+