Параметризация функции, включая строки в Spark Scala

Question

Вот мой набор данных:

case class CustomerStats(customerID: Integer, stat: Double)

Вот моя текущая функция:

def keepMaxStat(ds:Dataset[CustomerStats]) = {
    ds.groupByKey(_.customerID).
       reduceGroups( (x, y) => if (x.stat > y.stat) x else y )
}

Но на самом деле мой класс дел больше похож на:

case class CustomerStats(customerID: Integer, stat1: Double, stat2: Double)

Я пытаюсь параметризовать мою предыдущую функцию, например:

def keepMaxStat(ds:Dataset[CustomerStats], colName:String) = {
    ds.groupByKey(_.customerID).
       reduceGroups( (x, y) => if (x.colName > y.colName) x else y )
}

Но я получил эту ошибку:

ошибка: значение colName не является членом CustomerDateStats

Как я могу сделать это правильно?

Answer 1

Если вы хотите сохранить CustomerStats, вам придется управлять потоком на основе аргумента colName:

import scala.math.Ordering

def keepMaxStat(ds:Dataset[CustomerStats], colName:String) = {
  val ord = colName match {
    case "stat1" => Ordering.by[CustomerStats, Double](_.stat1)
    case "stat2" => Ordering.by[CustomerStats, Double](_.stat2)
  }
  ds.groupByKey(_.customerID).reduceGroups(
    ord.max(_, _)  
  )
}

или полностью переписать функцию, чтобы принять порядок:

def keepMaxStat(ds:Dataset[CustomerStats], ord: Ordering[CustomerStats]) = {
  ds.groupByKey(_.customerID).reduceGroups(
    ord.max(_, _)  
  )
}

keepMaxStat(ds, Ordering.by[CustomerStats, Double](_.stat1))

Answer 2

Если вы хотите сравнить Rows, вы должны сначала преобразовать Dataset[CustomerStats] в Dataset[Row], а затем использовать правильные get* методы доступа:

def keepMaxStat(ds:Dataset[CustomerStats], colName:String) = {
  ds.toDF.groupByKey(_.getAs[Integer]("customerID")).reduceGroups(
    (x, y) => if (x.getAs[Double](colName) > y.getAs[Double](colName)) x else y 
  )
}