Добавление новых столбцов с использованием forloop в искровом фрейме

Question

У меня есть искровой фрейм данных, который создается динамически. Есть также список столбцов, которые необходимо выбрать из фрейма данных.

Мне нужно перебрать список необходимых столбцов и проверить наличие этих столбцов в кадре данных. Если столбец данных найден, его нужно переименовать, в противном случае создайте новый столбец и поместите его со значением 'null'.

Я пытался использовать forloop, а затем, если условие, как показано ниже:

В моем фрейме данных есть столбцы: a.col1, a.col2, a.col3, a.col4, b.col2

val cols_needed = "a.col1,a.col2,a.col3,a.col4,b.col1,b.col2".split(",")
for(c <- cols_needed)
{
  if(Try(df(c)).isFailure)
  {
    df.withColumn(c, lit(null))
  }
  else
  {
    df.withColumn(`c`,df(c))
  }
}

Мне нужно использовать этот кадр данных после цикла для дальнейших операций.

Также пытался использовать цикл for с yield, но он дает мне массив данных.

Как обновить этот фрейм данных внутри forloop и, если условия. Потому что, когда я пытаюсь распечатать Schema за пределами цикла, у него старая схема, а не обновленная.

Любые предложения будут полезны.

Спасибо

Answer 1

API является неизменным, поэтому withColumn не изменяет существующее значение df, а возвращает новое.Например

val newDf = cols_needed.foldLeft(df) {
  case (tmpdf, c) =>
    if (Try(df(c)).isFailure) {
      tmpdf.withColumn(c, lit(null))
    } else {
      tmpdf.withColumn(`c`, df(c))
    }
}
// now newDf contains the added columns