Scala Partition/Collect Usage

落爺英雄遲暮 提交于 2019-12-02 16:37:04

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!

Adam Rabung

Not sure how to do it with collect without using mutable lists, but partition can use pattern matching as well (just a little more verbose)

List("a", 1, 2, "b", 19).partition { 
  case s:String => true
  case _ => false 
}

The signature of the normally-used collect on, say, Seq, is

collect[B](pf: PartialFunction[A,B]): Seq[B]

which is really a particular case of

collect[B, That](pf: PartialFunction[A,B])(
  implicit bf: CanBuildFrom[Seq[A], B, That]
): That

So if you use it in default mode, the answer is no, assuredly not: you get exactly one sequence out from it. If you follow CanBuildFrom through Builder, you see that it would be possible to make That actually be two sequences, but it would have no way of being told which sequence an item should go into, since the partial function can only say "yes, I belong" or "no, I do not belong".

So what do you do if you want to have multiple conditions that result in your list being split into a bunch of different pieces? One way is to create an indicator function A => Int, where your A is mapped into a numbered class, and then use groupBy. For example:

def optionClass(a: Any) = a match {
  case None => 0
  case Some(x) => 1
  case _ => 2
}
scala> List(None,3,Some(2),5,None).groupBy(optionClass)
res11: scala.collection.immutable.Map[Int,List[Any]] = 
  Map((2,List(3, 5)), (1,List(Some(2))), (0,List(None, None)))

Now you can look up your sub-lists by class (0, 1, and 2 in this case). Unfortunately, if you want to ignore some inputs, you still have to put them in a class (e.g. you probably don't care about the multiple copies of None in this case).

I use this. One nice thing about it is it combines partitioning and mapping in one iteration. One drawback is that it does allocate a bunch of temporary objects (the Either.Left and Either.Right instances)

/**
 * Splits the input list into a list of B's and a list of C's, depending on which type of value the mapper function returns.
 */
def mapSplit[A,B,C](in: List[A])(mapper: (A) => Either[B,C]): (List[B], List[C]) = {
  @tailrec
  def mapSplit0(in: List[A], bs: List[B], cs: List[C]): (List[B], List[C]) = {
    in match {
      case a :: as =>
        mapper(a) match {
          case Left(b)  => mapSplit0(as, b :: bs, cs     )
          case Right(c) => mapSplit0(as, bs,      c :: cs)
        }
      case Nil =>
        (bs.reverse, cs.reverse)
    }
  }

  mapSplit0(in, Nil, Nil)
}

val got = mapSplit(List(1,2,3,4,5)) {
  case x if x % 2 == 0 => Left(x)
  case y               => Right(y.toString * y)
}

assertEquals((List(2,4),List("1","333","55555")), got)

I could not find a satisfying solution to this basic problem here. I don't need a lecture on collect and don't care if this is someone's homework. Also, I don't want something that works only for List.

So here is my stab at it. Efficient and compatible with any TraversableOnce, even strings:

implicit class TraversableOnceHelper[A,Repr](private val repr: Repr)(implicit isTrav: Repr => TraversableOnce[A]) {

  def collectPartition[B,Left](pf: PartialFunction[A, B])
  (implicit bfLeft: CanBuildFrom[Repr, B, Left], bfRight: CanBuildFrom[Repr, A, Repr]): (Left, Repr) = {
    val left = bfLeft(repr)
    val right = bfRight(repr)
    val it = repr.toIterator
    while (it.hasNext) {
      val next = it.next
      if (!pf.runWith(left += _)(next)) right += next
    }
    left.result -> right.result
  }

  def mapSplit[B,C,Left,Right](f: A => Either[B,C])
  (implicit bfLeft: CanBuildFrom[Repr, B, Left], bfRight: CanBuildFrom[Repr, C, Right]): (Left, Right) = {
    val left = bfLeft(repr)
    val right = bfRight(repr)
    val it = repr.toIterator
    while (it.hasNext) {
      f(it.next) match {
        case Left(next) => left += next
        case Right(next) => right += next
      }
    }
    left.result -> right.result
  }
}

Example usages:

val (syms, ints) =
  Seq(Left('ok), Right(42), Right(666), Left('ko), Right(-1)) mapSplit identity

val ctx = Map('a -> 1, 'b -> 2) map {case(n,v) => n->(n,v)}
val (bound, unbound) = Vector('a, 'a, 'c, 'b) collectPartition ctx
println(bound: Vector[(Symbol, Int)], unbound: Vector[Symbol])

Starting Scala 2.13, most collections are now provided with a partitionMap method which partitions elements based on a function which returns either Right or Left.

That allows us to pattern match based on the type (which as a collect enables having specific types in the partitioned lists) or any other pattern:

 val (strings, ints) =
   List("a", 1, 2, "b", 19).partitionMap {
     case s: String => Left(s)
     case x: Int    => Right(x)
   }
 // strings: List[String] = List("a", "b")
 // ints: List[Int] = List(1, 2, 19)

Something like this could help

def partitionMap[IN, A, B](seq: Seq[IN])(function: IN => Either[A, B]): (Seq[A], Seq[B]) = {
  val (eitherLeft, eitherRight) = seq.map(function).partition(_.isLeft)
  eitherLeft.map(_.left.get) -> eitherRight.map(_.right.get)
}

To call it

val seq: Seq[Any] = Seq(1, "A", 2, "B")
val (ints, strings) = CollectionUtils.partitionMap(seq) {
  case int: Int    => Left(int)
  case str: String => Right(str)
}
ints shouldBe Seq(1, 2)
strings shouldBe Seq("A", "B")

Advantage is a simple API, similar with the one from Scala 2.12

Disadvantage; collection is ran twice and missing support for CanBuildFrom

标签
易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!