reduce by key on DS in Spark

reduceByKey.scala

def reduceBy[A, B](merge: (A, A) => A)(
    by: A => B
  )(ds: Dataset[A])(implicit session: SparkSession, encoderA: Encoder[A], encoderB: Encoder[B]): Dataset[A] = {

    def reducePartition(iter: Iterator[A]): Iterator[A] = {
      iter.toList
        .groupBy(by)
        .mapValues(values => values.reduce(merge))
        .values
        .toIterator
    }

    ds.mapPartitions(reducePartition)
      .groupByKey(by)
      .reduceGroups(merge)
      .map { case (_, updates) => updates }

  }