如何在spark scala中根据匹配条件从列名中提取数字？

Question

我有一个像这样的数据框架，

  val df = Seq((12,14,56,22,18),(11,54,83,91,12)).toDF("individual_rule_1_professional",
"individual_rule_2_amateur","team_rule_3_professional","team_rule_4_amateur","total")

我想把以个体开头和以业余开头的列的规则编号提取到一个数组中。因此，在本例中，它将返回2，因为在dataframe中只有一列以个体开始，以业余结束，且该列的规则编号为2。

我该怎么做呢？

非常感谢！祝你今天过得愉快。

Answer 1

可以使用df.columns提取列名

df.columns
// Array[String] = Array(individual_rule_1_professional, individual_rule_2_amateur, team_rule_3_professional, team_rule_4_amateur, total)

然后对其应用基本模式匹配以提取规则编号(我假设它始终是一个数字)

df.columns.filter(_.matches("individual_.*_amateur")).map(_.replaceAll("[^\\d]", ""))
// Array[String] = Array(2)

Answer 2

以下是另一种解决方案：

df.columns.filter(c => c.startsWith("individual") && c.endsWith("amateur"))
  .flatMap(_.split('_').flatMap(_.toIntOption))
  //or if the number you want is always in the second position
  //.flatMap(_.split('_')(2).toIntOption)
//res1: Array[Int] = Array(2)

与@phlantorvert的答案的区别：

如果单词“Array[String]

• Avoids”后面没有紧跟_字符(但仍然假设分隔符是_字符)，则
• 将输出Array[Int]而不是_任何问题。这更符合您指定的列名："...start with“和"...end with"