Scala Partition/Collect Usage

Adrian Modliszewski picture Adrian Modliszewski · Jan 24, 2011 · Viewed 46.3k times · Source

Is it possible to use one call to collect to make 2 new lists? If not, how can I do this using partition?

Answer

Kevin Wright picture Kevin Wright · Jan 24, 2011

collect (defined on TraversableLike and available in all subclasses) works with a collection and a PartialFunction. It also just so happens that a bunch of case clauses defined inside braces are a partial function (See section 8.5 of the Scala Language Specification [warning - PDF])

As in exception handling:

try {
  ... do something risky ...
} catch {
  //The contents of this catch block are a partial function
  case e: IOException => ...
  case e: OtherException => ...
}

It's a handy way to define a function that will only accept some values of a given type.

Consider using it on a list of mixed values:

val mixedList = List("a", 1, 2, "b", 19, 42.0) //this is a List[Any]
val results = mixedList collect {
  case s: String => "String:" + s
  case i: Int => "Int:" + i.toString
}

The argument to to collect method is a PartialFunction[Any,String]. PartialFunction because it's not defined for all possible inputs of type Any (that being the type of the List) and String because that's what all the clauses return.

If you tried to use map instead of collect, the the double value at the end of mixedList would cause a MatchError. Using collect just discards this, as well as any other value for which the PartialFunction is not defined.

One possible use would be to apply different logic to elements of the list:

var strings = List.empty[String]
var ints = List.empty[Int]
mixedList collect {
  case s: String => strings :+= s
  case i: Int => ints :+= i
}

Although this is just an example, using mutable variables like this is considered by many to be a war crime - So please don't do it!

A much better solution is to use collect twice:

val strings = mixedList collect { case s: String => s }
val ints = mixedList collect { case i: Int => i }

Or if you know for certain that the list only contains two types of values, you can use partition, which splits a collections into values depending on whether or not they match some predicate:

//if the list only contains Strings and Ints:
val (strings, ints) = mixedList partition { case s: String => true; case _ => false }

The catch here is that both strings and ints are of type List[Any], though you can easily coerce them back to something more typesafe (perhaps by using collect...)

If you already have a type-safe collection and want to split on some other property of the elements, then things are a bit easier for you:

val intList = List(2,7,9,1,6,5,8,2,4,6,2,9,8)
val (big,small) = intList partition (_ > 5)
//big and small are both now List[Int]s

Hope that sums up how the two methods can help you out here!