Scatter gather flatMap/groupBy fixup

tasks/src/main/scala/dagr/tasks/ScatterGather.scala

@@ -115,6 +115,17 @@ object ScatterGather {
     }
   }
 
+  /** Implementation of a Partitioner that takes the partitions from an existing Partitioner, and then groups them. */
+  private class GroupByPartitioner[Result, Key](partitioner: Partitioner[Result], f: Result => Key) extends SimpleInJvmTask with Partitioner[(Key, Seq[Result])] {
+    var partitions: Option[Seq[(Key, Seq[Result])]] = None
+    override def run(): Unit = {
+      partitioner.partitions match {
+        case None             => throw new IllegalStateException(s"partitioner.partitions called before partitions populated.")
+        case Some(_partition) => this.partitions = Some(_partition.groupBy(f).toSeq)
+      }
+    }
+  }
+
   /**
   * Implementation of a Scatter that is just a thinly veiled wrapper around the Scatterer
   * being used to generate the set of scatters/partitions to operate on.
@@ -126,8 +137,11 @@ object ScatterGather {
     override def gather[NextResult <: Task](f: Seq[Result] => NextResult): Gather[Result,NextResult] =
       throw new UnsupportedOperationException("gather not supported on an unmapped Scatter")
 
-    override def groupBy[Key](f: Result => Key) : Scatter[(Key, Seq[Result])] =
-      throw new UnsupportedOperationException("groupBy not supported on an unmapped Scatter")
+    override def groupBy[Key](f: Result => Key) : Scatter[(Key, Seq[Result])] = {
+      val grouper = new GroupByPartitioner[Result, Key](partitioner, f)
+      this ==> grouper.scatter
+      grouper.scatter
+    }
 
     override def flatMap[NextResult](f: Result => Scatter[NextResult]) : Scatter[NextResult] = {
       this.map(f).flatMap(identity)

tasks/src/test/scala/dagr/tasks/ScatterGatherTests.scala

@@ -107,6 +107,10 @@ class ScatterGatherTests extends UnitSpec with LazyLogging with BeforeAndAfterAl
     def run(): Unit = Io.writeLines(output, Seq(number.toString))
   }
 
+  private case class WriteNumberTuple(numbers: (Int, Int), output: Path) extends SimpleInJvmTask {
+    def run(): Unit = Io.writeLines(output, Seq(Seq(numbers._1, numbers._2).map(_.toString).mkString("\t")))
+  }
+
   "ScatterGather" should "run a simple scatter-gather pipeline on files" in {
     val lines = Seq("one", "one two", "one two three", "one two three four", "one two three four five")
     val lengths = Seq(1,2,3,4,5)
@@ -132,11 +136,7 @@ class ScatterGatherTests extends UnitSpec with LazyLogging with BeforeAndAfterAl
 
     val taskManager = buildTaskManager
     taskManager.addTask(pipeline)
-    taskManager.runToCompletion(true).foreach { case (task, info) =>
-      if (TaskStatus.isTaskNotDone(info.status)) {
-        println(s"${task.name} $info")
-      }
-    }
+    taskManager.runToCompletion(true)
 
     val sum1 = Io.readLines(sumOfCounts).next().toInt
     val sum2 = Io.readLines(sumOfSquares).next().toInt
@@ -351,4 +351,48 @@ class ScatterGatherTests extends UnitSpec with LazyLogging with BeforeAndAfterAl
 
     sumFlatMap shouldBe lengths.sum
   }
+
+  it should "flatMap a scatter, then groupBy and map" in {
+    val lines = Seq("one", "one two", "one two three", "one two three four", "one two three four five")
+
+    // setup the input and output
+    val input = tmp()
+    val countsByWordLengthOut = tmp()
+    Io.writeLines(input, lines)
+
+    val pipeline = new Pipeline() {
+      override def build(): Unit = {
+        // the initial scatter: scatters across lines
+        val scatter: Scatter[Path] = Scatter(SplitByLine(input=input))
+
+        // scatter from flatMap: each line is scattered across words, then flatMap makes a scatter across all words (all lines)
+        val scatterByWordFlatMap: Scatter[Path] = scatter.flatMap { pathToLine =>
+          val scatter = Scatter(SplitLineByWord(pathToLine))
+          root ==> scatter
+          scatter
+        }
+
+        // group by word length
+        val groupedByWordLength = scatterByWordFlatMap.groupBy { pathToWord => Io.readLines(pathToWord).next().length }
+
+        // map: count how many occurrences of words of a given length
+        val countsByWordLength = groupedByWordLength.map { case (wordLength, tasks) => WriteNumberTuple(numbers=(wordLength, tasks.length), output=tmp()) }
+
+        // gather: concatenate them all
+        countsByWordLength.gather { tasks => Concat(inputs = tasks.map(_.output), output = countsByWordLengthOut) }
+
+        root ==> scatter
+      }
+    }
+
+    val taskManager = buildTaskManager
+    taskManager.addTask(pipeline)
+    taskManager.runToCompletion(true)
+
+    val outLines = Io.readLines(countsByWordLengthOut).toList
+    outLines.map { line =>
+      val Array(wordLength: String, count: String) = line.split("\t")
+      (wordLength.toInt, count.toInt)
+    }.sortBy(_._1) should contain theSameElementsInOrderAs Seq((3, 9), (4, 3), (5, 3))
+  }
 }