groupBy on scala iterator

IteratorUtils.scala

package utils

import scala.collection.mutable
import scala.language.implicitConversions

object IteratorUtils {

  /**
    * Implicit class that adds "groupBy" functionality to an iterator. The result is a mutable.HashMap and will be
    * stored in memory in its entirety
    * @param iterator the iterator to operate on
    * @tparam ElementType type of iterator elements
    */
  implicit class IteratorGroupBy[ElementType](iterator: Iterator[ElementType]) {
    /**
      * groups each element in iterator by the key function and stores them in a vector
      * @param makeKey the function to get the key from an element
      * @tparam KeyType self expl.
      * @example 
      *  case class Toy(id: Int, price: Int)
      *  val itCollect = Iterator(Toy(1,100),Toy(1,200),Toy(2,200),Toy(3,10),Toy(1,2234))
      *  itCollect.groupByCollect(_.id)
      *  _______________________________________
      *  res0: scala.collection.mutable.HashMap[Int,Vector[Toy]] = Map(2 -> Vector(Toy(2,200)), 1 -> Vector(Toy(1,2234), Toy(1,200), Toy(1,100)), 3 -> Vector(Toy(3,10)))
      */
    def groupByCollect[KeyType](makeKey: (ElementType) => KeyType): mutable.HashMap[KeyType, Vector[ElementType]] = {
      val internal = new mutable.HashMap[KeyType, Vector[ElementType]]()
      iterator.foreach(elementA => {
        val key = makeKey(elementA)
        val res = internal.get(key)
        internal(key) = res.map(value => elementA +: value).getOrElse(Vector(elementA))
      })
      internal
    }

    /**
      * Keeps only one element per key from iterator. compare function is used to merge/compare/fold two different elements with the same key
      * @param makeKey the function to get the key from an element
      * @param compare given two elements gives you the element to store
      * @tparam KeyType self expl.
      * @example
      *  case class Toy(id: Int, price: Int)
      *  val itFold = Iterator(Toy(1,100),Toy(1,200),Toy(2,200),Toy(3,10),Toy(1,2234))
      *  itFold.groupByFold(_.id,(a,b) => if(a.price < b.price) a else b)
      *  _______________________________________
      *  res0: scala.collection.mutable.HashMap[Int,Toy] = Map(2 -> Toy(2,200), 1 -> Toy(1,100), 3 -> Toy(3,10))
      */
    def groupByFold[KeyType](makeKey: (ElementType) => KeyType, compare: (ElementType,ElementType) => ElementType): mutable.HashMap[KeyType, ElementType] = {
      val internal = new mutable.HashMap[KeyType, ElementType]()
      iterator.foreach(elementA => {
        val key = makeKey(elementA)
        val res = internal.get(key)
        internal(key) = res.map(value => compare(elementA,value)).getOrElse(elementA)
      })
      internal
    }
  }
}