Fix NoSuchElementException when a state is not set but timeoutThreshold is defined

zsxwing · zsxwing · commit 784db1b9d4c0 · 2016-02-04T11:02:33.000-08:00
diff --git a/streaming/src/main/scala/org/apache/spark/streaming/rdd/MapWithStateRDD.scala b/streaming/src/main/scala/org/apache/spark/streaming/rdd/MapWithStateRDD.scala
@@ -57,7 +57,8 @@ private[streaming] object MapWithStateRDDRecord {
       val returned = mappingFunction(batchTime, key, Some(value), wrappedState)
       if (wrappedState.isRemoved) {
         newStateMap.remove(key)
-      } else if (wrappedState.isUpdated || timeoutThresholdTime.isDefined) {
+      } else if (wrappedState.isUpdated
+          || (wrappedState.exists && timeoutThresholdTime.isDefined)) {
         newStateMap.put(key, wrappedState.get(), batchTime.milliseconds)
       }
       mappedData ++= returned
diff --git a/streaming/src/test/scala/org/apache/spark/streaming/rdd/MapWithStateRDDSuite.scala b/streaming/src/test/scala/org/apache/spark/streaming/rdd/MapWithStateRDDSuite.scala
@@ -190,6 +190,11 @@ class MapWithStateRDDSuite extends SparkFunSuite with RDDCheckpointTester with B
       timeoutThreshold = Some(initialTime + 1), removeTimedoutData = true,
       expectedStates = Nil, expectedTimingOutStates = Nil, expectedRemovedStates = Seq(123))
 
+    // If a state is not set but timeoutThreshold is defined, we should ignore this state.
+    // Previously it threw NoSuchElementException (SPARK-13195).
+    assertRecordUpdate(initStates = Seq(), data = Seq("noop"),
+      timeoutThreshold = Some(initialTime + 1), removeTimedoutData = true,
+      expectedStates = Nil, expectedTimingOutStates = Nil)
   }
 
   test("states generated by MapWithStateRDD") {
