Skip to content
Snippets Groups Projects
Commit 1638fcb0 authored by Tathagata Das's avatar Tathagata Das
Browse files

Fixed updateStateByKey to work with primitive types.

parent 131be5d6
No related branches found
No related tags found
No related merge requests found
...@@ -377,7 +377,7 @@ extends Serializable { ...@@ -377,7 +377,7 @@ extends Serializable {
* corresponding state key-value pair will be eliminated. * corresponding state key-value pair will be eliminated.
* @tparam S State type * @tparam S State type
*/ */
def updateStateByKey[S <: AnyRef : ClassManifest]( def updateStateByKey[S: ClassManifest](
updateFunc: (Seq[V], Option[S]) => Option[S] updateFunc: (Seq[V], Option[S]) => Option[S]
): DStream[(K, S)] = { ): DStream[(K, S)] = {
updateStateByKey(updateFunc, defaultPartitioner()) updateStateByKey(updateFunc, defaultPartitioner())
...@@ -392,7 +392,7 @@ extends Serializable { ...@@ -392,7 +392,7 @@ extends Serializable {
* @param numPartitions Number of partitions of each RDD in the new DStream. * @param numPartitions Number of partitions of each RDD in the new DStream.
* @tparam S State type * @tparam S State type
*/ */
def updateStateByKey[S <: AnyRef : ClassManifest]( def updateStateByKey[S: ClassManifest](
updateFunc: (Seq[V], Option[S]) => Option[S], updateFunc: (Seq[V], Option[S]) => Option[S],
numPartitions: Int numPartitions: Int
): DStream[(K, S)] = { ): DStream[(K, S)] = {
...@@ -408,7 +408,7 @@ extends Serializable { ...@@ -408,7 +408,7 @@ extends Serializable {
* @param partitioner Partitioner for controlling the partitioning of each RDD in the new DStream. * @param partitioner Partitioner for controlling the partitioning of each RDD in the new DStream.
* @tparam S State type * @tparam S State type
*/ */
def updateStateByKey[S <: AnyRef : ClassManifest]( def updateStateByKey[S: ClassManifest](
updateFunc: (Seq[V], Option[S]) => Option[S], updateFunc: (Seq[V], Option[S]) => Option[S],
partitioner: Partitioner partitioner: Partitioner
): DStream[(K, S)] = { ): DStream[(K, S)] = {
...@@ -431,7 +431,7 @@ extends Serializable { ...@@ -431,7 +431,7 @@ extends Serializable {
* @param rememberPartitioner Whether to remember the paritioner object in the generated RDDs. * @param rememberPartitioner Whether to remember the paritioner object in the generated RDDs.
* @tparam S State type * @tparam S State type
*/ */
def updateStateByKey[S <: AnyRef : ClassManifest]( def updateStateByKey[S: ClassManifest](
updateFunc: (Iterator[(K, Seq[V], Option[S])]) => Iterator[(K, S)], updateFunc: (Iterator[(K, Seq[V], Option[S])]) => Iterator[(K, S)],
partitioner: Partitioner, partitioner: Partitioner,
rememberPartitioner: Boolean rememberPartitioner: Boolean
......
...@@ -7,7 +7,7 @@ import spark.storage.StorageLevel ...@@ -7,7 +7,7 @@ import spark.storage.StorageLevel
import spark.streaming.{Duration, Time, DStream} import spark.streaming.{Duration, Time, DStream}
private[streaming] private[streaming]
class StateDStream[K: ClassManifest, V: ClassManifest, S <: AnyRef : ClassManifest]( class StateDStream[K: ClassManifest, V: ClassManifest, S: ClassManifest](
parent: DStream[(K, V)], parent: DStream[(K, V)],
updateFunc: (Iterator[(K, Seq[V], Option[S])]) => Iterator[(K, S)], updateFunc: (Iterator[(K, Seq[V], Option[S])]) => Iterator[(K, S)],
partitioner: Partitioner, partitioner: Partitioner,
......
...@@ -151,10 +151,10 @@ class BasicOperationsSuite extends TestSuiteBase { ...@@ -151,10 +151,10 @@ class BasicOperationsSuite extends TestSuiteBase {
) )
val updateStateOperation = (s: DStream[String]) => { val updateStateOperation = (s: DStream[String]) => {
val updateFunc = (values: Seq[Int], state: Option[RichInt]) => { val updateFunc = (values: Seq[Int], state: Option[Int]) => {
Some(new RichInt(values.foldLeft(0)(_ + _) + state.map(_.self).getOrElse(0))) Some(values.foldLeft(0)(_ + _) + state.getOrElse(0))
} }
s.map(x => (x, 1)).updateStateByKey[RichInt](updateFunc).map(t => (t._1, t._2.self)) s.map(x => (x, 1)).updateStateByKey[Int](updateFunc)
} }
testOperation(inputData, updateStateOperation, outputData, true) testOperation(inputData, updateStateOperation, outputData, true)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment