Snapshot a pipe in the REPL

scalding-core/src/main/scala/com/twitter/scalding/Source.scala

@@ -108,10 +108,9 @@ abstract class Source extends java.io.Serializable {
   }
 
   /**
-   * write the pipe and return the input so it can be chained into
-   * the next operation
+   * Write the given pipe and return the new pipe which was added as the tail
    */
-  def writeFrom(pipe: Pipe)(implicit flowDef: FlowDef, mode: Mode) = {
+  def writeFromAndGetTail(pipe: Pipe)(implicit flowDef: FlowDef, mode: Mode) = {
     checkFlowDefNotNull
 
     //insane workaround for scala compiler bug
@@ -124,7 +123,17 @@ abstract class Source extends java.io.Serializable {
       case (test: TestMode, false) => pipe
       case _ => transformForWrite(pipe)
     }
-    flowDef.addTail(new Pipe(sinkName, newPipe))
+    val finalPipe = new Pipe(sinkName, newPipe)
+    flowDef.addTail(finalPipe)
+    finalPipe
+  }
+
+  /**
+   * write the pipe but return the input so it can be chained into
+   * the next operation
+   */
+  def writeFrom(pipe: Pipe)(implicit flowDef: FlowDef, mode: Mode) = {
+    writeFromAndGetTail(pipe)
     pipe
   }
 

scalding-repl/src/main/scala/com/twitter/scalding/ReplImplicits.scala

@@ -17,6 +17,7 @@ package com.twitter.scalding
 
 import cascading.flow.FlowDef
 import cascading.pipe.Pipe
+import typed.KeyedListLike
 
 /**
  * Object containing various implicit conversions required to create Scalding flows in the REPL.
@@ -132,8 +133,18 @@ object ReplImplicits extends FieldConversions {
    * @return a ShellPipe wrapping the specified Pipe.
    */
   implicit def pipeToShellPipe(pipe: Pipe): ShellObj[Pipe] = new ShellObj(pipe)
-  implicit def typedPipeToShellPipe[T](pipe: TypedPipe[T]): ShellObj[TypedPipe[T]] =
-    new ShellObj(pipe)
-  implicit def keyedListToShellPipe[K, V](pipe: KeyedList[K, V]): ShellObj[KeyedList[K, V]] =
-    new ShellObj(pipe)
+
+  /**
+   * Convert KeyedListLike to enriched ShellTypedPipe
+   * (e.g. allows .snapshot to be called on Grouped, CoGrouped, etc)
+   */
+  implicit def keyedListLikeToShellTypedPipe[K, V, T[K, +V] <: KeyedListLike[K, V, T]](kll: KeyedListLike[K, V, T]) = new ShellTypedPipe(kll.toTypedPipe)
+
+  /**
+   * Enrich TypedPipe for the shell
+   * (e.g. allows .snapshot to be called on it)
+   */
+  implicit def typedPipeToShellTypedPipe[T](pipe: TypedPipe[T]): ShellTypedPipe[T] =
+    new ShellTypedPipe[T](pipe)
+
 }

scalding-repl/src/main/scala/com/twitter/scalding/ShellPipe.scala

@@ -16,7 +16,10 @@
 package com.twitter.scalding
 
 import cascading.flow.Flow
+import cascading.flow.FlowDef
 import cascading.pipe.Pipe
+import java.util.UUID
+import com.twitter.scalding.ReplImplicits._
 
 /**
  * Adds ability to run a pipe in the REPL.
@@ -30,12 +33,12 @@ class ShellObj[T](obj: T) {
    * @param args that should be used to construct the job.
    * @return a job that can be used to run the data pipeline.
    */
-  private[scalding] def getJob(args: Args, inmode: Mode): Job = new Job(args) {
+  private[scalding] def getJob(args: Args, inmode: Mode, inFlowDef: FlowDef): Job = new Job(args) {
     /**
      *  The flow definition used by this job, which should be the same as that used by the user
      *  when creating their pipe.
      */
-    override val flowDef = ReplImplicits.flowDef
+    override val flowDef = inFlowDef
 
     override def mode = inmode
 
@@ -89,24 +92,92 @@ class ShellObj[T](obj: T) {
      */
     override def buildFlow: Flow[_] = {
       val flow = super.buildFlow
-      ReplImplicits.resetFlowDef()
       flow
     }
   }
 
   /**
    * Runs this pipe as a Scalding job.
    */
-  def run() {
-    val args = new Args(Map())
-    getJob(args, ReplImplicits.mode).run
-  }
+  def run(inFlowDef: FlowDef = ReplImplicits.flowDef) =
+    getJob(new Args(Map()), ReplImplicits.mode, inFlowDef).run
 
   def toList[R](implicit ev: T <:< TypedPipe[R], manifest: Manifest[R]): List[R] = {
     import ReplImplicits._
     ev(obj).toPipe("el").write(Tsv("item"))
     run()
     TypedTsv[R]("item").toIterator.toList
   }
+
 }
 
+/**
+ * Enrichment on TypedPipes allowing them to be run locally, independent of the overall flow.
+ * @param pipe to wrap
+ */
+class ShellTypedPipe[T](pipe: TypedPipe[T]) extends ShellObj[TypedPipe[T]](pipe) {
+
+  /**
+   * Iterator for all pipes reachable from this pipe (recursively using 'Pipe.getPrevious')
+   */
+  def upstreamPipes(inpipe: Pipe): Iterator[Pipe] =
+    Iterator
+      .iterate(Seq(inpipe))(pipes => for (pipe <- pipes; prev <- pipe.getPrevious) yield prev)
+      .takeWhile(_.length > 0)
+      .flatten
+
+  /**
+   * Construct a new FlowDef for only the flow that ends with the given pipe.
+   * That is, it copies over only the sources and sinks that contribute to the
+   * flow, allowing repl users to build up flows incrementally.
+   */
+  def localizedFlow(tailPipe: Pipe): FlowDef = {
+    val newFlow = getEmptyFlowDef
+
+    val sourceTaps = flowDef.getSources
+    val newSrcs = newFlow.getSources
+
+    upstreamPipes(tailPipe)
+      .filter(_.getParent == null)
+      .flatMap(_.getHeads)
+      .foreach(head =>
+        if (!newSrcs.containsKey(head.getName))
+          newFlow.addSource(head, sourceTaps.get(head.getName)))
+
+    newFlow.addTailSink(tailPipe, flowDef.getSinks.get(tailPipe.getName))
+
+    newFlow
+  }
+
+  /**
+   * Shorthand for .write(dest).run
+   */
+  def save(dest: TypedSink[T] with Mappable[T]): TypedPipe[T] = {
+
+    val d = dest
+    val thisPipe = pipe.toPipe(d.sinkFields)(d.setter)
+    val outPipe = d.writeFromAndGetTail(thisPipe)
+
+    run(localizedFlow(outPipe))
+
+    TypedPipe.from(d)
+  }
+
+  /**
+   * Save snapshot of a typed pipe to a temporary sequence file.
+   * @return A TypedPipe to a new Source, reading from the sequence file.
+   */
+  def snapshot: TypedPipe[T] = {
+    import ReplImplicits._
+
+    // come up with unique temporary filename
+    // TODO: refactor into TemporarySequenceFile class
+    val tmpSeq = "/tmp/scalding-repl/snapshot-" + UUID.randomUUID() + ".seq"
+    val outPipe = SequenceFile(tmpSeq, 'record).writeFromAndGetTail(pipe.toPipe('record))
+
+    run(localizedFlow(outPipe))
+
+    TypedPipe.fromSingleField[T](SequenceFile(tmpSeq))
+  }
+
+}

scalding-repl/src/test/scala/com/twitter/scalding/ReplTest.scala

@@ -0,0 +1,44 @@
+/*
+Copyright 2012 Twitter, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package com.twitter.scalding
+
+object ReplTest {
+  import ReplImplicits._
+
+  def test() {
+    val hello = TypedPipe.from(TextLine("tutorial/data/hello.txt"))
+
+    val wordScores =
+      TypedPipe.from(OffsetTextLine("tutorial/data/words.txt"))
+        .map{ case (offset, word) => (word, offset) }
+        .group
+
+    // snapshot intermediate results without wiring everything up
+    val s1 = hello.snapshot
+
+    val s2 = hello.save(TypedTsv("dump.tsv"))
+
+    // use snapshot in further flows
+    val linesByWord = s1.flatMap(_.split("\\s+")).groupBy(_.toLowerCase)
+    val counts = linesByWord.size
+
+    // ensure snapshot enrichment works on KeyedListLike (CoGrouped, UnsortedGrouped), too
+    val s3 = counts.snapshot
+    val s4 = linesByWord.join(wordScores).snapshot
+  }
+
+}