Scala method to side effect on map and return it

Question

What is the best way to apply a function to each element of a Map and at the end return the same Map, unchanged, so that it can be used in further operations?

I'd like to avoid:

myMap.map(el => {
  effectfullFn(el)
  el
})

to achieve syntax like this:

myMap
  .mapEffectOnKV(effectfullFn)
  .foreach(println)

map is not what I'm looking for, because I have to specify what comes out of the map (as in the first code snippet), and I don't want to do that.

I want a special operation that knows/assumes that the map elements should be returned without change after the side-effect function has been executed.

In fact, this would be so useful to me, I'd like to have it for Map, Array, List, Seq, Iterable... The general idea is to peek at the elements to do something, then automatically return these elements.

The real case I'm working on looks like this:

 calculateStatistics(trainingData, indexMapLoaders)
   .superMap { (featureShardId, shardStats) =>
      val outputDir = summarizationOutputDir + "/" + featureShardId
      val indexMap = indexMapLoaders(featureShardId).indexMapForDriver()
      IOUtils.writeBasicStatistics(sc, shardStats, outputDir, indexMap)
    }

Once I have calculated the statistics for each shard, I want to append the side effect of saving them to disk, and then just return those statistics, without having to create a val and having that val's name be the last statement in the function, e.g.:

val stats = calculateStatistics(trainingData, indexMapLoaders)
stats.foreach { (featureShardId, shardStats) =>
  val outputDir = summarizationOutputDir + "/" + featureShardId
  val indexMap = indexMapLoaders(featureShardId).indexMapForDriver()
  IOUtils.writeBasicStatistics(sc, shardStats, outputDir, indexMap)
}
stats

It's probably not very hard to implement, but I was wondering if there was something in Scala already for that.

Answer 1

Function cannot be effectful by definition, so I wouldn't expect anything convenient in scala-lib. However, you can write a wrapper:

def tap[T](effect: T => Unit)(x: T) = {
  effect(x)
  x
}

Example:

scala> Map(1 -> 1, 2 -> 2)
         .map(tap(el => el._1 + 5 -> el._2))
         .foreach(println)
(1,1)
(2,2)

You can also define an implicit:

implicit class TapMap[K,V](m: Map[K,V]){
  def tap(effect: ((K,V)) => Unit): Map[K,V] = m.map{x =>
    effect(x)
    x
  }
}

Examples:

scala> Map(1 -> 1, 2 -> 2).tap(el => el._1 + 5 -> el._2).foreach(println)
(1,1)
(2,2)

To abstract more, you can define this implicit on TraversableOnce, so it would be applicable to List, Set and so on if you need it:

implicit class TapTraversable[Coll[_], T](m: Coll[T])(implicit ev: Coll[T] <:< TraversableOnce[T]){
  def tap(effect: T => Unit): Coll[T] = {
    ev(m).foreach(effect)
    m
  }
}

scala> List(1,2,3).tap(println).map(_ + 1)
1
2
3
res24: List[Int] = List(2, 3, 4)

scala> Map(1 -> 1).tap(println).toMap //`toMap` is needed here for same reasons as it needed when you do `.map(f).toMap`
(1,1)
res5: scala.collection.immutable.Map[Int,Int] = Map(1 -> 1)

scala> Set(1).tap(println)
1
res6: scala.collection.immutable.Set[Int] = Set(1)

It's more useful, but requires some "mamba-jumbo" with types, as Coll[_] <: TraversableOnce[_] doesn't work (Scala 2.12.1), so I had to use an evidence for that.

You can also try CanBuildFrom approach: How to enrich a TraversableOnce with my own generic map?

---

Overall recommendation about dealing with passthrough side-effects on iterators is to use Streams (scalaz/fs2/monix) and Task, so they've got an observe (or some analog of it) function that does what you want in async (if needed) way.

---

My answer before you provided example of what you want

You can represent effectful computation without side-effects and have distinct values that represent state before and after:

scala> val withoutSideEffect = Map(1 -> 1, 2 -> 2)
withoutSideEffect: scala.collection.immutable.Map[Int,Int] = Map(1 -> 1, 2 -> 2)                                                                       

scala> val withSideEffect = withoutSideEffect.map(el => el._1 + 5 -> (el._2 + 5))
withSideEffect: scala.collection.immutable.Map[Int,Int] = Map(6 -> 6, 7 -> 7)

scala> withoutSideEffect //unchanged
res0: scala.collection.immutable.Map[Int,Int] = Map(1 -> 1, 2 -> 2)

scala> withSideEffect //changed
res1: scala.collection.immutable.Map[Int,Int] = Map(6 -> 6, 7 -> 7)