sparkr sparkr - 1 month ago 23
Scala Question

Scala Pattern Matching on Generic Traits

I have a trait that is defined like this:

sealed trait MyTrait[+A] {
def name: String
def id: Long
def someType: A
}


I now have some case classes that extend this trait as below:

object MyTrait {
case class One[+A](name: String, id: Long, someField: String, someType: A) extends MyTrait[A]

case class Two[+A](name: String, id: Long, someField: String, someType: A) extends MyTrait[A]

case class Three[+A](name: String, id: Long, someType: A) extends MyTrait[A]
}


Now, I have a List that might contain one of these case classes and I have to collect the individual case classes from the given List and I use instanceOf checks as below:

val myBigList = List(One[SomeType], Two[SomeType], Three[SomeType], One[SomeType])

val allOnes = myBigList.collect {
case elem if elem.isInstanceOf[One[_]] => elem.asInstanceOf[One]]
}

val allTwos = myBigList.collect {
case elem if elem.isInstanceOf[Two[_]] => elem.asInstanceOf[Two]]
}

val allThrees = myBigList.collect {
case elem if elem.isInstanceOf[Three[_]] => elem.asInstanceOf[Three]]
}


Is there a better way? I mean is there a way to avoid the isInstanceOf and asInstanceOf?

Answer

This is good if you want to One[String] and One[Int] into one bucket meaning if you do not care about inner types.

But you can do this in one single pass instead of three passes using foldLeft.

def bucketize(list: List[MyTrait[_]]): (List[One[_]], List[Two[_]], List[Three[_]]) = {
  list.foldLeft(((List.empty[One[_]], List.empty[Two[_]], List.empty[Three[_]]))){ (r, c) => 
    val ((one, two, three)) = r
    c match { 
      case x: One[_] => ((one ++ List(x), two, three))
      case x: Two[_] => ((one, two ++ List(x), three))
      case x: Three[_] => ((one, two, three ++ List(x)))
    }
   }
}

You get triplet of lists, first list is List[One], second is List[Two] and so on ...

usage:

val ((one, two, three))  = bucketize(list)