Как мне извлечь числа из имен столбцов на основе условия соответствия в spark scala?

#scala #apache-spark

#scala #apache-spark

Вопрос:

У меня есть такой фрейм данных,

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

и я хотел бы извлечь номера правил столбцов, которые начинаются с individual и заканчиваются amateur, в массив. Таким образом, в этом случае он вернет [2] , потому что во фрейме данных есть только один столбец, который начинается с individual и заканчивается amateur , а номер правила для этого столбца равен 2 .

Как мне это сделать?

Большое спасибо! Хорошего дня.

Ответ №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)
  

Ответ №2:

Вот альтернативное решение:

 df.columns.filter(c => c.startsWith("individual") amp;amp; 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[Int] вместо Array[String]
  • Позволяет избежать любых проблем, если за словом «индивидуальный» не сразу следует _ символ (но все же предполагается, что разделитель является _ символом). Это лучше соответствует вашей спецификации, что имя столбца должно: «… начинаться с» и «… заканчиваться»