Spark - как добавить новые столбцы в исходный кадр данных после разделения?

Question

Я хотел разделить столбец DataFrame (col3) на несколько столбцов (col_a, col_b, col_c).И заменить один столбец несколькими столбцами.

+---------------------------------+
|col1|col2| col3  |col4|col5|col6 |
|   a|   b| a,b,c | d  |  n | r   | 
|   d|   e| f,f,f | p  |  b | null|
+---------------------------------+ 

Поэтому я разделил столбец следующим образом:

val new_DF = original_DF.withColumn("_tmp", split($"col3", "\\,")).select(
  $"_tmp".getItem(0).as("col_a"),
  $"_tmp".getItem(1).as("col_b"),
  $"_tmp".getItem(2).as("col_c")
).drop("_tmp")

Но new_DF содержит только новые три столбца,Как я могу сделать новые столбцы частью оригинальной original_DF без col3?

UPD: Мне удалось достичь желаемого результата, но это не выглядит красивосовсем.Буду благодарен за любые предложения по улучшению этого подхода.

val new_DF = original_DF
  .withColumn("col_a", split($"col3", "\\,").getItem(0))
  .withColumn("col_b", split($"col3", "\\,").getItem(1))
  .withColumn("col_c", split($"col3", "\\,").getItem(2))
  .drop("col3")

Answer 1

Это должно сработать:

val new_DF = original_DF.withColumn("_tmp", split($"col3", "\\,")).select(
  $"_tmp".getItem(0).as("col_a"),
  $"_tmp".getItem(1).as("col_b"),
  $"_tmp".getItem(2).as("col_c"),
  original_DF.columns.filter(_ != "col3"): _*)
)

Выбор всех столбцов из original_df без col3

Answer 2

без дубликатов:

val withArrayColumn = original_DF.withColumn("col3Array", split($"col3", ","))
val columnNames = Seq("a", "b", "c").zipWithIndex
val result = columnNames
  .foldLeft(withArrayColumn)((updated, columnIndex) => updated.withColumn("col_" + columnIndex._1, $"col3Array".getItem(columnIndex._2)))
  .drop("col3Array", "col3")

Answer 3

Это небольшое улучшение

val new_DF = original_DF.select(col("col_1"),
          col("col_2"),
          col("col3").getItem(0).as("col_a"),
          col("col3").getItem(1).as("col_b"),
          col("col3").getItem(2).as("col_c"),
          col("col_4"),
          col("col_5"),
          col("col_6"))