Addressed type issue in StateSpec.function

Author: tdas

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

@@ -150,15 +150,15 @@ object StateSpec {
    * @tparam StateType    Class of the states data
    * @tparam EmittedType  Class of the emitted data
    */
-  def function[ValueType, StateType, EmittedType](
+  def function[KeyType, ValueType, StateType, EmittedType](
       trackingFunction: (Option[ValueType], State[StateType]) => EmittedType
-    ): StateSpec[Any, ValueType, StateType, EmittedType] = {
+    ): StateSpec[KeyType, 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)
+    new StateSpecImpl(wrappedFunction)
   }
 }
 

streaming/src/main/scala/org/apache/spark/streaming/dstream/PairDStreamFunctions.scala

@@ -380,7 +380,8 @@ class PairDStreamFunctions[K, V](self: DStream[(K, V)])
    */
   @Experimental
   def trackStateByKey[StateType: ClassTag, EmittedType: ClassTag](
-    spec: StateSpec[K, V, StateType, EmittedType]): TrackStateDStream[K, StateType, EmittedType] = {
+      spec: StateSpec[K, V, StateType, EmittedType]
+    ): TrackStateDStream[K, V, StateType, EmittedType] = {
     new TrackStateDStreamImpl[K, V, StateType, EmittedType](
       self,
       spec.asInstanceOf[StateSpecImpl[K, V, StateType, EmittedType]]

streaming/src/main/scala/org/apache/spark/streaming/dstream/TrackStateDStream.scala

@@ -39,7 +39,7 @@ import org.apache.spark.streaming.rdd.{TrackStateRDD, TrackStateRDDRecord}
  * @tparam EmittedType Class of the emitted records
  */
 @Experimental
-sealed abstract class TrackStateDStream[KeyType, StateType, EmittedType: ClassTag](
+sealed abstract class TrackStateDStream[KeyType, ValueType, StateType, EmittedType: ClassTag](
     ssc: StreamingContext) extends DStream[EmittedType](ssc) {
 
   /** Return a pair DStream where each RDD is the snapshot of the state of all the keys. */
@@ -51,7 +51,7 @@ private[streaming] class TrackStateDStreamImpl[
     KeyType: ClassTag, ValueType: ClassTag, StateType: ClassTag, EmittedType: ClassTag](
     dataStream: DStream[(KeyType, ValueType)],
     spec: StateSpecImpl[KeyType, ValueType, StateType, EmittedType])
-  extends TrackStateDStream[KeyType, StateType, EmittedType](dataStream.context) {
+  extends TrackStateDStream[KeyType, ValueType, StateType, EmittedType](dataStream.context) {
 
   private val internalStream =
     new InternalTrackStateDStream[KeyType, ValueType, StateType, EmittedType](dataStream, spec)
@@ -78,6 +78,14 @@ private[streaming] class TrackStateDStreamImpl[
     internalStream.flatMap {
       _.stateMap.getAll().map { case (k, s, _) => (k, s) }.toTraversable }
   }
+
+  def keyClass: Class[_] = implicitly[ClassTag[KeyType]].runtimeClass
+
+  def valueClass: Class[_] = implicitly[ClassTag[ValueType]].runtimeClass
+
+  def stateClass: Class[_] = implicitly[ClassTag[StateType]].runtimeClass
+
+  def emittedClass: Class[_] = implicitly[ClassTag[EmittedType]].runtimeClass
 }
 
 /**

streaming/src/test/scala/org/apache/spark/streaming/TrackStateByKeySuite.scala

@@ -24,6 +24,7 @@ import scala.reflect.ClassTag
 
 import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll}
 
+import org.apache.spark.streaming.dstream.{TrackStateDStream, TrackStateDStreamImpl}
 import org.apache.spark.util.{ManualClock, Utils}
 import org.apache.spark.{SparkConf, SparkContext, SparkFunSuite}
 
@@ -166,7 +167,8 @@ class TrackStateByKeySuite extends SparkFunSuite with BeforeAndAfterAll with Bef
       sum
     }
 
-    testOperation(inputData, StateSpec.function(trackStateFunc), outputData, stateData)
+    testOperation[String, Int, Int](
+      inputData, StateSpec.function(trackStateFunc), outputData, stateData)
   }
 
   test("trackStateByKey - basic operations with advanced API") {
@@ -213,6 +215,55 @@ class TrackStateByKeySuite extends SparkFunSuite with BeforeAndAfterAll with Bef
     testOperation(inputData, StateSpec.function(trackStateFunc), outputData, stateData)
   }
 
+  test("trackStateByKey - type inferencing and class tags") {
+
+    // Simple track state function with value as Int, state as Double and emitted type as Double
+    val simpleFunc = (value: Option[Int], state: State[Double]) => {
+      0L
+    }
+
+    // Advanced track state function with key as String, value as Int, state as Double and
+    // emitted type as Double
+    val advancedFunc = (time: Time, key: String, value: Option[Int], state: State[Double]) => {
+      Some(0L)
+    }
+
+    def testTypes(dstream: TrackStateDStream[_, _, _, _]): Unit = {
+      val dstreamImpl = dstream.asInstanceOf[TrackStateDStreamImpl[_, _, _, _]]
+      assert(dstreamImpl.keyClass === classOf[String])
+      assert(dstreamImpl.valueClass === classOf[Int])
+      assert(dstreamImpl.stateClass === classOf[Double])
+      assert(dstreamImpl.emittedClass === classOf[Long])
+    }
+
+    val inputStream = new TestInputStream[(String, Int)](ssc, Seq.empty, numPartitions = 2)
+
+    // Defining StateSpec inline with trackStateByKey and simple function implicitly gets the types
+    val simpleFunctionStateStream1 = inputStream.trackStateByKey(
+      StateSpec.function(simpleFunc).numPartitions(1))
+    testTypes(simpleFunctionStateStream1)
+
+    // Separately defining StateSpec with simple function requires explicitly specifying types
+    val simpleFuncSpec = StateSpec.function[String, Int, Double, Long](simpleFunc)
+    val simpleFunctionStateStream2 = inputStream.trackStateByKey(simpleFuncSpec)
+    testTypes(simpleFunctionStateStream2)
+
+    // Separately defining StateSpec with advanced function implicitly gets the types
+    val advFuncSpec1 = StateSpec.function(advancedFunc)
+    val advFunctionStateStream1 = inputStream.trackStateByKey(advFuncSpec1)
+    testTypes(advFunctionStateStream1)
+
+    // Defining StateSpec inline with trackStateByKey and advanced func implicitly gets the types
+    val advFunctionStateStream2 = inputStream.trackStateByKey(
+      StateSpec.function(simpleFunc).numPartitions(1))
+    testTypes(advFunctionStateStream2)
+
+    // Defining StateSpec inline with trackStateByKey and advanced func implicitly gets the types
+    val advFuncSpec2 = StateSpec.function[String, Int, Double, Long](advancedFunc)
+    val advFunctionStateStream3 = inputStream.trackStateByKey[Double, Long](advFuncSpec2)
+    testTypes(advFunctionStateStream3)
+  }
+
   test("trackStateByKey - states as emitted records") {
     val inputData =
       Seq(