Используйте два списка переменных для добавления столбцов scala

Question

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

Seq one - это что-то вроде: Seq("red", "blue", "green", "yellow", "violet")

, а Seq two - что-то вроде: Seq("child", "teen", "adult", "senior")

У меня также есть столбец, представляющий собой строку в формате: s "$ color + $ age-score = $ score", содержащую каждую комбинацию цветов и возрастов, с полученным результатом, поэтому20 различных возрастных показателей.

В настоящее время я делаю что-то вроде

finalDF.withColumn("red_child", getScore("red", "child"))
.withColumn("red_teen", getScore("red", "teen"))
.withColumn("red_adult", getScore("red", "adult"))

и т. Д. Для всех 20 возможных комбинаций, причем getScore является вспомогательной функцией, которая заботится оregex.

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

Спасибо.

Answer 1

Вы можете просто select дополнительные столбцы, полученные из списка кортежей, сгенерированного с помощью for-comprehension, как показано ниже:

val colors = Seq("red", "blue", "green", "yellow", "violet")
val ageGroups = Seq("child", "teen", "adult", "senior")

val colPairs = for { c <- colors; a <- ageGroups } yield (c, a)

def getScore(c: String, a: String): Column = ???

df.select( df.columns.map(col) ++ colPairs.map{ case (c, a) =>
    getScore(c, a).as(c + "_" + a)
  }: _*
)

В качестве альтернативы, используйте foldLeft, чтобы просмотреть список colPairs и добавить столбцы с помощью withColumn:

colPairs.foldLeft(df){ case (accDF, (c, a)) =>
  accDF.withColumn(c + "_" + a, getScore(c, a))
}