Snapshot a pipe in the REPL

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

@@ -16,6 +16,7 @@ limitations under the License.
 package com.twitter.scalding
 
 import cascading.pipe.Pipe
+import cascading.flow.FlowDef
 
 /**
  * This object has all the implicit functions and values that are used
@@ -29,4 +30,10 @@ import cascading.pipe.Pipe
  */
 object Dsl extends FieldConversions with java.io.Serializable {
   implicit def pipeToRichPipe(pipe: Pipe): RichPipe = new RichPipe(pipe)
+
+  /**
+   * Enrichment on FlowDef
+   */
+  implicit def flowDefToRichFlowDef(fd: FlowDef): RichFlowDef = new RichFlowDef(fd)
+
 }

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

@@ -0,0 +1,101 @@
+/*
+Copyright 2014 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
+
+import cascading.flow.FlowDef
+import cascading.pipe.Pipe
+
+/**
+ * This is an enrichment-pattern class for cascading.flow.FlowDef.
+ * The rule is to never use this class directly in input or return types, but
+ * only to add methods to FlowDef.
+ */
+class RichFlowDef(val fd: FlowDef) {
+  // RichPipe and RichFlowDef implicits
+  import Dsl._
+
+  def copy: FlowDef = {
+    val newFd = new FlowDef
+    newFd.merge(fd)
+    newFd
+  }
+
+  private[scalding] def mergeMisc(o: FlowDef) {
+    fd.addTags(o.getTags)
+    fd.addTraps(o.getTraps)
+    fd.addCheckpoints(o.getCheckpoints)
+    fd.setAssertionLevel(o.getAssertionLevel)
+    fd.setName(o.getName)
+    fd
+  }
+
+  /**
+   * Mutate current flow def to add all sources/sinks/etc from given FlowDef
+   */
+  def merge(o: FlowDef): FlowDef = {
+    fd.addSources(o.getSources)
+    fd.addSinks(o.getSinks)
+    fd.addTails(o.getTails)
+    fd.mergeMisc(o)
+    fd
+  }
+
+  /**
+   * New flow def with only sources upstream from tails.
+   */
+  def withoutUnusedSources: FlowDef = {
+    import collection.JavaConverters._
+
+    // find all heads reachable from the tails (as a set of names)
+    val heads = fd.getTails.asScala
+      .flatMap(_.getHeads).map(_.getName).toSet
+    // add taps associated with heads to localFlow
+    val filteredSources = fd.getSources.asScala.filterKeys(src => heads.contains(src)).asJava
+
+    val newFd = fd.copy
+    newFd.getSources.clear()
+    newFd.addSources(filteredSources)
+
+    newFd
+  }
+
+  /**
+   * FlowDef that only includes things upstream from the given Pipe
+   */
+  def onlyUpstreamFrom(pipe: Pipe): FlowDef = {
+    val newFd = new FlowDef
+    // don't copy any sources/sinks
+    newFd.mergeMisc(fd)
+
+    val sourceTaps = fd.getSources
+    val newSrcs = newFd.getSources
+
+    pipe.upstreamPipes
+      .filter(_.getPrevious.length == 0) // implies _ is a head
+      .foreach { head =>
+        if (!newSrcs.containsKey(head.getName))
+          newFd.addSource(head, sourceTaps.get(head.getName))
+      }
+
+    val sinks = fd.getSinks
+    if (sinks.containsKey(pipe.getName)) {
+      newFd.addTailSink(pipe, sinks.get(pipe.getName))
+    }
+
+    newFd
+  }
+
+}

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

@@ -645,6 +645,17 @@ class RichPipe(val pipe: Pipe) extends java.io.Serializable with JoinAlgorithms
     val setter = unpacker.newSetter(toFields)
     pipe.mapTo(fields) { input: T => input } (conv, setter)
   }
+
+  /**
+   * Set of pipes reachable from this pipe (transitive closure of 'Pipe.getPrevious')
+   */
+  def upstreamPipes: Set[Pipe] =
+    Iterator
+      .iterate(Seq(pipe))(pipes => for (p <- pipes; prev <- p.getPrevious) yield prev)
+      .takeWhile(_.length > 0)
+      .flatten
+      .toSet
+
 }
 
 /**

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

@@ -108,14 +108,14 @@ abstract class Source extends java.io.Serializable {
   }
 
   /**
-   * write the pipe and return the input so it can be chained into
+   * 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) = {
+  def writeFrom(pipe: Pipe)(implicit flowDef: FlowDef, mode: Mode): Pipe = {
     checkFlowDefNotNull
 
     //insane workaround for scala compiler bug
-    val sinks = flowDef.getSinks().asInstanceOf[JMap[String, Any]]
+    val sinks = flowDef.getSinks.asInstanceOf[JMap[String, Any]]
     val sinkName = this.toString
     if (!sinks.containsKey(sinkName)) {
       sinks.put(sinkName, createTap(Write)(mode))
@@ -124,7 +124,8 @@ abstract class Source extends java.io.Serializable {
       case (test: TestMode, false) => pipe
       case _ => transformForWrite(pipe)
     }
-    flowDef.addTail(new Pipe(sinkName, newPipe))
+    val outPipe = new Pipe(sinkName, newPipe)
+    flowDef.addTail(outPipe)
     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.
@@ -46,6 +47,72 @@ object ReplImplicits extends FieldConversions {
     fd
   }
 
+  /**
+   * Gets a job that can be used to run the data pipeline.
+   *
+   * @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, 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 = inFlowDef
+
+    override def mode = inmode
+
+    /**
+     * Obtains a configuration used when running the job.
+     *
+     * This overridden method uses the same configuration as a standard Scalding job,
+     * but adds options specific to KijiScalding, including adding a jar containing compiled REPL
+     * code to the distributed cache if the REPL is running.
+     *
+     * @return the configuration that should be used to run the job.
+     */
+    override def config: Map[AnyRef, AnyRef] = {
+      // Use the configuration from Scalding Job as our base.
+      val configuration: Map[AnyRef, AnyRef] = super.config
+
+      /** Appends a comma to the end of a string. */
+      def appendComma(str: Any): String = str.toString + ","
+
+      // If the REPL is running, we should add tmpjars passed in from the command line,
+      // and a jar of REPL code, to the distributed cache of jobs run through the REPL.
+      val replCodeJar = ScaldingShell.createReplCodeJar()
+      val tmpJarsConfig: Map[String, String] =
+        if (replCodeJar.isDefined) {
+          Map("tmpjars" -> {
+            // Use tmpjars already in the configuration.
+            configuration
+              .get("tmpjars")
+              .map(appendComma)
+              .getOrElse("") +
+              // And a jar of code compiled by the REPL.
+              "file://" + replCodeJar.get.getAbsolutePath
+          })
+        } else {
+          // No need to add the tmpjars to the configuration
+          Map()
+        }
+
+      configuration ++ tmpJarsConfig
+    }
+
+  }
+
+  /**
+   * Runs this pipe as a Scalding job.
+   *
+   * Automatically cleans up the flowDef to include only sources upstream from tails.
+   */
+  def run(implicit flowDef: FlowDef) = {
+    import Dsl.flowDefToRichFlowDef
+
+    getJob(new Args(Map()), ReplImplicits.mode, flowDef.withoutUnusedSources).run
+  }
+
   /**
    * Converts a Cascading Pipe to a Scalding RichPipe. This method permits implicit conversions from
    * Pipe to RichPipe.
@@ -125,15 +192,16 @@ object ReplImplicits extends FieldConversions {
   }
 
   /**
-   * Converts a Cascading Pipe to a Scalding ShellPipe. This method permits implicit conversions
-   * from Pipe to ShellPipe.
-   *
-   * @param pipe to convert to a ShellPipe.
-   * @return a ShellPipe wrapping the specified Pipe.
+   * Convert KeyedListLike to enriched ShellTypedPipe
+   * (e.g. allows .snapshot to be called on Grouped, CoGrouped, etc)
    */
-  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)
+  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)
+
 }