Search code examples
scalafor-comprehensionflatmap

Performance difference in flatMap compared to for-comprehension


I am still new to Scala, and one of the things I read about is that for-comprehension is equivalent to a flatMap to a certain extent. However, in my code (below), flatMap is taking at least as twice as long to compute. What could be the reason for this?

This is the slow one:

facts.flatMap(f => factActionsMap(f)).filter(_.isValid(facts))

This is the fast equivalent one:

for {
  f <- facts
  a <- factActionsMap(f)
  if a.isValid(facts)
} yield a

factActionsMap is a map between Fact and a Set[Action]. facts is just a Set[Fact].


Solution

  • Let's check the translation:

    scala> trait Fact
    defined trait Fact
    
    scala> trait Action { def isValid(s: Set[Fact]): Boolean }
    defined trait Action
    
    scala> def facts: Set[Fact] = ???
    facts: Set[Fact]
    
    scala> def factActionsMap: Map[Fact, Set[Action]] = ???
    factActionsMap: Map[Fact,Set[Action]]
    
    scala> import scala.reflect.runtime.{universe => u}
    import scala.reflect.runtime.{universe=>u}
    
    scala> val expr = u reify {
         | for {
         |   f <- facts
         |   a <- factActionsMap(f)
         |   if a.isValid(facts)
         | } yield a
         | }
    expr: reflect.runtime.universe.Expr[scala.collection.immutable.Set[Action]] = Expr[scala.collection.immutable.Set[Action]]($read.f
    acts.flatMap(((f) => $read.factActionsMap.apply(f).withFilter(((a) => a.isValid($read.facts))).map(((a) => a))(Set.canBuildFrom)))
    (Set.canBuildFrom))
    
    scala> u show expr.tree
    res0: String = $read.facts.flatMap(((f) => $read.factActionsMap.apply(f).withFilter(((a) => a.isValid($read.facts))).map(((a) => a
    ))(Set.canBuildFrom)))(Set.canBuildFrom)
    

    So, removing REPL stuff (everything starting with $) and the implicit parameters, plus reformatting, we get:

    facts.flatMap(f => factActionsMap(f).withFilter(a => a.isValid(facts)).map(a => a))
    

    There are two main differences between this and what you came up with. First, withFilter is applied to fastActionsMap(f) result, whereas you apply to facts.flatMap result. This means flatMap will work over all results, instead of just the ones accepted.

    Second, it uses withFilter instead of filter, which avoids creating an extra collection.