REPL: Add toIterator (and related methods)

project/Build.scala

@@ -282,6 +282,11 @@ object ScaldingBuild extends Build {
   ).settings(
     name := "scalding-repl",
     previousArtifact := None,
+    initialCommands in console := """
+      import com.twitter.scalding._
+      import com.twitter.scalding.ReplImplicits._
+      import com.twitter.scalding.ReplImplicitContext._
+      """,
     libraryDependencies <++= (scalaVersion) { scalaVersion => Seq(
       "org.scala-lang" % "jline" % scalaVersion,
       "org.scala-lang" % "scala-compiler" % scalaVersion,

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

@@ -120,6 +120,8 @@ trait HadoopMode extends Mode {
   override def openForRead(tap: Tap[_, _, _]) = {
     val htap = tap.asInstanceOf[Tap[JobConf, _, _]]
     val conf = new JobConf(jobConf)
+    // copy over Config defaults
+    Config.default.toMap.foreach{ case (k, v) => conf.set(k, v) }
     val fp = new HadoopFlowProcess(conf)
     htap.retrieveSourceFields(fp)
     htap.sourceConfInit(fp, conf)

scalding-core/src/main/scala/com/twitter/scalding/source/TypedSequenceFile.scala

@@ -0,0 +1,37 @@
+/*
+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.source
+
+import cascading.tuple.Fields
+import com.twitter.scalding._
+import com.twitter.scalding.SequenceFile
+
+/**
+ * SequenceFile with explicit types. Useful for debugging flows using the Typed API.
+ * Not to be used for permanent storage: uses Kryo serialization which may not be
+ * consistent across JVM instances. Use Thrift sources instead.
+ */
+class TypedSequenceFile[T](path: String) extends SequenceFile(path, Fields.FIRST) with Mappable[T] with TypedSink[T] {
+  override def converter[U >: T] =
+    TupleConverter.asSuperConverter[T, U](TupleConverter.singleConverter[T])
+  override def setter[U <: T] =
+    TupleSetter.asSubSetter[T, U](TupleSetter.singleSetter[T])
+}
+
+object TypedSequenceFile {
+  def apply[T](path: String): TypedSequenceFile[T] = new TypedSequenceFile[T](path)
+}

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

@@ -25,10 +25,11 @@ import scala.util.{ Failure, Success }
  * Most of these conversions come from the [[com.twitter.scalding.Job]] class.
  */
 object ReplImplicits extends FieldConversions {
+
   /** Implicit flowDef for this Scalding shell session. */
-  implicit var flowDef: FlowDef = getEmptyFlowDef
+  var flowDef: FlowDef = getEmptyFlowDef
   /** Defaults to running in local mode if no mode is specified. */
-  implicit var mode: Mode = com.twitter.scalding.Local(false)
+  var mode: Mode = com.twitter.scalding.Local(false)
 
   /**
    * Sets the flow definition in implicit scope to an empty flow definition.
@@ -53,8 +54,7 @@ object ReplImplicits extends FieldConversions {
    *
    * Automatically cleans up the flowDef to include only sources upstream from tails.
    */
-  def run(implicit flowDef: FlowDef): Option[JobStats] = {
-    import Dsl.flowDefToRichFlowDef
+  def run(implicit fd: FlowDef, md: Mode): Option[JobStats] = {
 
     def config = {
       val conf = Config.default
@@ -81,7 +81,7 @@ object ReplImplicits extends FieldConversions {
     }
 
     // TODO: This is not getting any UniqueID, so counters will not work with REPL
-    ExecutionContext.newContext(config, flowDef, mode).waitFor match {
+    ExecutionContext.newContext(config, fd, md).waitFor match {
       case Success(stats) => Some(stats)
       case Failure(e) =>
         println("Flow execution failed!")
@@ -149,7 +149,7 @@ object ReplImplicits extends FieldConversions {
    */
   implicit def iterableToPipe[T](
     iterable: Iterable[T])(implicit setter: TupleSetter[T],
-      converter: TupleConverter[T]): Pipe = {
+      converter: TupleConverter[T], fd: FlowDef, md: Mode): Pipe = {
     iterableToSource(iterable)(setter, converter).read
   }
 
@@ -164,8 +164,8 @@ object ReplImplicits extends FieldConversions {
    */
   implicit def iterableToRichPipe[T](
     iterable: Iterable[T])(implicit setter: TupleSetter[T],
-      converter: TupleConverter[T]): RichPipe = {
-    RichPipe(iterableToPipe(iterable)(setter, converter))
+      converter: TupleConverter[T], fd: FlowDef, md: Mode): RichPipe = {
+    RichPipe(iterableToPipe(iterable)(setter, converter, fd, md))
   }
 
   /**
@@ -181,4 +181,22 @@ object ReplImplicits extends FieldConversions {
   implicit def typedPipeToShellTypedPipe[T](pipe: TypedPipe[T]): ShellTypedPipe[T] =
     new ShellTypedPipe[T](pipe)
 
+  /**
+   * Enrich ValuePipe for the shell
+   * (e.g. allows .toOption to be called on it)
+   */
+  implicit def valuePipeToShellValuePipe[T](pipe: ValuePipe[T]): ShellValuePipe[T] =
+    new ShellValuePipe[T](pipe)
+
+}
+
+/**
+ * Implicit FlowDef and Mode, import in the REPL to have the global context implicitly
+ * used everywhere.
+ */
+object ReplImplicitContext {
+  /** Implicit flowDef for this Scalding shell session. */
+  implicit var fd = ReplImplicits.flowDef
+  /** Defaults to running in local mode if no mode is specified. */
+  implicit var md = ReplImplicits.mode
 }

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

@@ -60,7 +60,10 @@ class ScaldingILoop
 
   addThunk {
     intp.beQuietDuring {
-      intp.addImports("com.twitter.scalding._", "com.twitter.scalding.ReplImplicits._")
+      intp.addImports(
+        "com.twitter.scalding._",
+        "com.twitter.scalding.ReplImplicits._",
+        "com.twitter.scalding.ReplImplicitContext._")
     }
   }
 }

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

@@ -15,54 +15,105 @@
 
 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._
+
+import cascading.flow.FlowDef
+import cascading.tuple.Fields
+import com.twitter.scalding.typed._
+import scala.collection.JavaConverters._
+import com.twitter.scalding.source.TypedSequenceFile
 
 /**
  * Enrichment on TypedPipes allowing them to be run locally, independent of the overall flow.
  * @param pipe to wrap
  */
 class ShellTypedPipe[T](pipe: TypedPipe[T]) {
   import Dsl.flowDefToRichFlowDef
+  import ReplImplicits._
 
   /**
    * Shorthand for .write(dest).run
    */
-  def save(dest: TypedSink[T] with Mappable[T]): TypedPipe[T] = {
+  def save(dest: TypedSink[T] with Mappable[T])(implicit fd: FlowDef, md: Mode): TypedPipe[T] = {
 
     val p = pipe.toPipe(dest.sinkFields)(dest.setter)
 
-    val localFlow = flowDef.onlyUpstreamFrom(p)
-    dest.writeFrom(p)(localFlow, mode)
-    run(localFlow)
+    val localFlow = fd.onlyUpstreamFrom(p)
+    dest.writeFrom(p)(localFlow, md)
+    run(localFlow, md)
 
-    TypedPipe.from(dest)
+    TypedPipe.from(dest)(fd, md)
   }
 
   /**
    * 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] = {
-
-    // come up with unique temporary filename
-    // TODO: refactor into TemporarySequenceFile class
-    val tmpSeq = "/tmp/scalding-repl/snapshot-" + UUID.randomUUID() + ".seq"
-    val dest = SequenceFile(tmpSeq, 'record)
-    val p = pipe.toPipe('record)
-
-    val localFlow = flowDef.onlyUpstreamFrom(p)
-    dest.writeFrom(p)(localFlow, mode)
-    run(localFlow)
+  def snapshot(implicit fd: FlowDef, md: Mode): TypedPipe[T] = {
+    val p = pipe.toPipe(0)
+    val localFlow = fd.onlyUpstreamFrom(p)
+    md match {
+      case _: CascadingLocal => // Local or Test mode
+        val dest = new MemorySink[T]
+        dest.writeFrom(p)(localFlow, md)
+        run(localFlow, md)
+        TypedPipe.from(dest.readResults)(fd, md)
+      case _: HadoopMode =>
+        // come up with unique temporary filename
+        // TODO: refactor into TemporarySequenceFile class
+        val tmpSeq = "/tmp/scalding-repl/snapshot-" + UUID.randomUUID + ".seq"
+        val dest = TypedSequenceFile[T](tmpSeq)
+        dest.writeFrom(p)(localFlow, md)
+        run(localFlow, md)
+        TypedPipe.from(dest)(fd, md)
+    }
+  }
 
-    TypedPipe.fromSingleField[T](SequenceFile(tmpSeq))
+  /**
+   * Create a (local) iterator over the pipe. For non-trivial pipes (anything except
+   * a head-pipe reading from a source), a snapshot is automatically created and
+   * iterated over.
+   * @return local iterator
+   */
+  def toIterator(implicit fd: FlowDef, md: Mode): Iterator[T] = pipe match {
+    // if this is just a Converter on a head pipe
+    // (true for the first pipe on a source, e.g. a snapshot pipe)
+    case TypedPipeInst(p, fields, Converter(conv)) if p.getPrevious.isEmpty =>
+      val srcs = fd.getSources
+      if (srcs.containsKey(p.getName)) {
+        val tap = srcs.get(p.getName)
+        md.openForRead(tap).asScala.map(tup => conv(tup.selectEntry(fields)))
+      } else {
+        sys.error("Invalid head: pipe has no previous, but there is no registered source.")
+      }
+    // if it's already just a wrapped iterable (MemorySink), just return it
+    case IterablePipe(iter, _, _) => iter.toIterator
+    // otherwise, snapshot the pipe and get an iterator on that
+    case _ =>
+      pipe.snapshot.toIterator
   }
 
-  // TODO: add back `toList` based on `snapshot` this time
+  /**
+   * Create a list from the pipe in memory. Uses `ShellTypedPipe.toIterator`.
+   * Warning: user must ensure that the results will actually fit in memory.
+   */
+  def toList(implicit fd: FlowDef, md: Mode): List[T] = toIterator.toList
 
-  // TODO: add `dump` to view contents without reading into memory
+  /**
+   * Print the contents of a pipe to stdout. Uses `ShellTypedPipe.toIterator`.
+   */
+  def dump(implicit fd: FlowDef, md: Mode): Unit = toIterator.foreach(println(_))
 
 }
+
+class ShellValuePipe[T](vp: ValuePipe[T]) {
+  import ReplImplicits.typedPipeToShellTypedPipe
+  def toOption(implicit fd: FlowDef, md: Mode): Option[T] = vp match {
+    case _: EmptyValue => None
+    case LiteralValue(v) => Some(v)
+    case ComputedValue(tp) => tp.snapshot.toList match {
+      case Nil => None
+      case v :: Nil => Some(v)
+    }
+  }
+}