Question

我想使用两个Seq将列添加到数据框。

以下之一是这样的： Seq("red", "blue", "green", "yellow", "violet")

和Seq两个类似： 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是负责正则表达式的辅助函数。

由于我使用withColumn 20次，因此使代码很难阅读。我想知道是否有任何方法可以使用颜色和年龄的两个Seq循环并向数据框中添加列，从而使代码看起来更简洁。

谢谢。

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))
}

使用两个变量列表添加scala列

1 个答案: