Addressed PR comments

Author: tdas

streaming/src/main/scala/org/apache/spark/streaming/State.scala

@@ -29,31 +29,23 @@ import org.apache.spark.annotation.Experimental
  *
  * Scala example of using `State`:
  * {{{
- *    def trackStateFunc(key: String, data: Option[Int], wrappedState: State[Int]): Option[Int] = {
- *
+ *    // A tracking function that maintains an integer state and return a String
+ *    def trackStateFunc(data: Option[Int], state: State[Int]): Option[String] = {
  *      // Check if state exists
  *      if (state.exists) {
- *
- *        val existingState = wrappedState.get  // Get the existing state
- *
- *        val shouldRemove = ...      // Decide whether to remove the state
- *
+ *        val existingState = state.get  // Get the existing state
+ *        val shouldRemove = ...         // Decide whether to remove the state
  *        if (shouldRemove) {
- *
- *          wrappedState.remove()     // Remove the state
- *
+ *          state.remove()     // Remove the state
  *        } else {
- *
  *          val newState = ...
- *          wrappedState(newState)    // Set the new state
- *
+ *          state.update(newState)    // Set the new state
  *        }
  *      } else {
- *
  *        val initialState = ...
  *        state.update(initialState)  // Set the initial state
- *
  *      }
+ *      ... // return something
  *    }
  *
  * }}}
@@ -98,7 +90,7 @@ sealed abstract class State[S] {
 
   /**
    * Whether the state is timing out and going to be removed by the system after the current batch.
-   * This timeou can occur if timeout duration has been specified in the
+   * This timeout can occur if timeout duration has been specified in the
    * [[org.apache.spark.streaming.StateSpec StatSpec]] and the key has not received any new data
    * for that timeout duration.
    */
@@ -114,16 +106,11 @@ sealed abstract class State[S] {
   }
 }
 
-private[streaming]
-object State {
-  implicit def toOption[S](state: State[S]): Option[S] = state.getOption()
-}
-
 /** Internal implementation of the [[State]] interface */
 private[streaming] class StateImpl[S] extends State[S] {
 
   private var state: S = null.asInstanceOf[S]
-  private var defined: Boolean = true
+  private var defined: Boolean = false
   private var timingOut: Boolean = false
   private var updated: Boolean = false
   private var removed: Boolean = false
@@ -134,13 +121,18 @@ private[streaming] class StateImpl[S] extends State[S] {
   }
 
   override def get(): S = {
-    state
+    if (defined) {
+      state
+    } else {
+      throw new NoSuchElementException("State is not set")
+    }
   }
 
   override def update(newState: S): Unit = {
     require(!removed, "Cannot update the state after it has been removed")
     require(!timingOut, "Cannot update the state that is timing out")
     state = newState
+    defined = true
     updated = true
   }
 
@@ -151,6 +143,8 @@ private[streaming] class StateImpl[S] extends State[S] {
   override def remove(): Unit = {
     require(!timingOut, "Cannot remove the state that is timing out")
     require(!removed, "Cannot remove the state that has already been removed")
+    defined = false
+    updated = false
     removed = true
   }
 
@@ -167,7 +161,7 @@ private[streaming] class StateImpl[S] extends State[S] {
   }
 
   /**
-   * Internal method to update the state data and reset internal flags in `this`.
+   * Update the internal data and flags in `this` to the given state option.
    * This method allows `this` object to be reused across many state records.
    */
   def wrap(optionalState: Option[S]): Unit = {
@@ -186,7 +180,7 @@ private[streaming] class StateImpl[S] extends State[S] {
   }
 
   /**
-   * Internal method to update the state data and reset internal flags in `this`.
+   * Update the internal data and flags in `this` to the given state that is going to be timed out.
    * This method allows `this` object to be reused across many state records.
    */
   def wrapTiminoutState(newState: S): Unit = {

streaming/src/main/scala/org/apache/spark/streaming/StateSpec.scala

@@ -22,6 +22,7 @@ import scala.reflect.ClassTag
 import org.apache.spark.annotation.Experimental
 import org.apache.spark.api.java.JavaPairRDD
 import org.apache.spark.rdd.RDD
+import org.apache.spark.util.ClosureCleaner
 import org.apache.spark.{HashPartitioner, Partitioner}
 
 
@@ -37,28 +38,33 @@ import org.apache.spark.{HashPartitioner, Partitioner}
  *
  * Example in Scala:
  * {{{
- *    val spec = StateSpec(trackingFunction).numPartitions(10)
+ *    def trackingFunction(data: Option[ValueType], wrappedState: State[StateType]): EmittedType = {
+ *      ...
+ *    }
+ *
+ *    val spec = StateSpec.function(trackingFunction).numPartitions(10)
  *
  *    val emittedRecordDStream = keyValueDStream.trackStateByKey[StateType, EmittedDataType](spec)
  * }}}
  *
  * Example in Java:
  * {{{
- *    StateStateSpec[StateType, EmittedDataType] spec =
- *      StateStateSpec.create[StateType, EmittedDataType](trackingFunction).numPartition(10);
+ *    StateStateSpec[KeyType, ValueType, StateType, EmittedDataType] spec =
+ *      StateStateSpec.function[KeyType, ValueType, StateType, EmittedDataType](trackingFunction)
+ *                    .numPartition(10);
  *
  *    JavaDStream[EmittedDataType] emittedRecordDStream =
  *      javaPairDStream.trackStateByKey[StateType, EmittedDataType](spec);
  * }}}
  */
 @Experimental
-sealed abstract class StateSpec[K, V, S, T] extends Serializable {
+sealed abstract class StateSpec[KeyType, ValueType, StateType, EmittedType] extends Serializable {
 
   /** Set the RDD containing the initial states that will be used by `trackStateByKey`*/
-  def initialState(rdd: RDD[(K, S)]): this.type
+  def initialState(rdd: RDD[(KeyType, StateType)]): this.type
 
   /** Set the RDD containing the initial states that will be used by `trackStateByKey`*/
-  def initialState(javaPairRDD: JavaPairRDD[K, S]): this.type
+  def initialState(javaPairRDD: JavaPairRDD[KeyType, StateType]): this.type
 
   /**
    * Set the number of partitions by which the state RDDs generated by `trackStateByKey`
@@ -93,15 +99,20 @@ sealed abstract class StateSpec[K, V, S, T] extends Serializable {
  *
  * Example in Scala:
  * {{{
- *    val spec = StateSpec(trackingFunction).numPartitions(10)
+ *    def trackingFunction(data: Option[ValueType], wrappedState: State[StateType]): EmittedType = {
+ *      ...
+ *    }
+ *
+ *    val spec = StateSpec.function(trackingFunction).numPartitions(10)
  *
  *    val emittedRecordDStream = keyValueDStream.trackStateByKey[StateType, EmittedDataType](spec)
  * }}}
  *
  * Example in Java:
  * {{{
- *    StateStateSpec[StateType, EmittedDataType] spec =
- *      StateStateSpec.create[StateType, EmittedDataType](trackingFunction).numPartition(10);
+ *    StateStateSpec[KeyType, ValueType, StateType, EmittedDataType] spec =
+ *      StateStateSpec.function[KeyType, ValueType, StateType, EmittedDataType](trackingFunction)
+ *                    .numPartition(10);
  *
  *    JavaDStream[EmittedDataType] emittedRecordDStream =
  *      javaPairDStream.trackStateByKey[StateType, EmittedDataType](spec);
@@ -115,16 +126,17 @@ object StateSpec {
    * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]] (Scala) or a
    * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]] (Java).
    * @param trackingFunction The function applied on every data item to manage the associated state
-   *                         and generate the emitted data and
+   *                         and generate the emitted data
    * @tparam KeyType      Class of the keys
    * @tparam ValueType    Class of the values
    * @tparam StateType    Class of the states data
    * @tparam EmittedType  Class of the emitted data
    */
-  def apply[KeyType, ValueType, StateType, EmittedType](
-      trackingFunction: (KeyType, Option[ValueType], State[StateType]) => Option[EmittedType]
+  def function[KeyType, ValueType, StateType, EmittedType](
+      trackingFunction: (Time, KeyType, Option[ValueType], State[StateType]) => Option[EmittedType]
     ): StateSpec[KeyType, ValueType, StateType, EmittedType] = {
-    new StateSpecImpl[KeyType, ValueType, StateType, EmittedType](trackingFunction)
+    ClosureCleaner.clean(trackingFunction, checkSerializable = true)
+    new StateSpecImpl(trackingFunction)
   }
 
   /**
@@ -133,24 +145,28 @@ object StateSpec {
    * [[org.apache.spark.streaming.dstream.PairDStreamFunctions pair DStream]] (Scala) or a
    * [[org.apache.spark.streaming.api.java.JavaPairDStream JavaPairDStream]] (Java).
    * @param trackingFunction The function applied on every data item to manage the associated state
-   *                         and generate the emitted data and
-   * @tparam KeyType      Class of the keys
+   *                         and generate the emitted data
    * @tparam ValueType    Class of the values
    * @tparam StateType    Class of the states data
    * @tparam EmittedType  Class of the emitted data
    */
-  def create[KeyType, ValueType, StateType, EmittedType](
-      trackingFunction: (KeyType, Option[ValueType], State[StateType]) => Option[EmittedType]
-    ): StateSpec[KeyType, ValueType, StateType, EmittedType] = {
-    apply(trackingFunction)
+  def function[ValueType, StateType, EmittedType](
+      trackingFunction: (Option[ValueType], State[StateType]) => EmittedType
+    ): StateSpec[Any, ValueType, StateType, EmittedType] = {
+    ClosureCleaner.clean(trackingFunction, checkSerializable = true)
+    val wrappedFunction =
+      (time: Time, key: Any, value: Option[ValueType], state: State[StateType]) => {
+        Some(trackingFunction(value, state))
+      }
+    new StateSpecImpl[Any, ValueType, StateType, EmittedType](wrappedFunction)
   }
 }
 
 
 /** Internal implementation of [[org.apache.spark.streaming.StateSpec]] interface. */
 private[streaming]
 case class StateSpecImpl[K, V, S, T](
-    function: (K, Option[V], State[S]) => Option[T]) extends StateSpec[K, V, S, T] {
+    function: (Time, K, Option[V], State[S]) => Option[T]) extends StateSpec[K, V, S, T] {
 
   require(function != null)
 
@@ -186,7 +202,7 @@ case class StateSpecImpl[K, V, S, T](
 
   // ================= Private Methods =================
 
-  private[streaming] def getFunction(): (K, Option[V], State[S]) => Option[T] = function
+  private[streaming] def getFunction(): (Time, K, Option[V], State[S]) => Option[T] = function
 
   private[streaming] def getInitialStateRDD(): Option[RDD[(K, S)]] = Option(initialStateRDD)