Передать строку в UDF и выбрать столбец на основе сопоставления с образцом

Question

Как я могу добиться следующего, передав строку в udf?

val df1 = df.withColumn("col_Z", 
              when($"col_x" === "a", $"col_A")
              .when($"col_x" === "b", $"col_B")
              .when($"col_x" === "c", $"col_C")
              .when($"col_x" === "d", $"col_D")
              .when($"col_x" === "e", $"col_E")
              .when($"col_x" === "f", $"col_F")
              .when($"col_x" === "g", $"col_G")
      )

Насколько я понимаю, только столбцы могут передаваться в качестве аргументов UDF в Scala Spark.

Я взглянул на этот вопрос:

Как передать всю строку в UDF - фильтр Spark DataFrame

и попытался реализовать этот udf:

def myUDF(r:Row) = udf {

 val z : Float = r.getAs("col_x") match {
      case "a" => r.getAs("col_A")
      case "b" => r.getAs("col_B")
      case other => lit(0.0)
   }
 z
}

но я получаю ошибку несоответствия типов:

 error: type mismatch;
 found   : String("a")
 required: Nothing
 case "a" => r.getAs("col_A")
      ^

Что я делаю не так?