[SPARK-19168][STRUCTURED STREAMING] StateStore should be aborted upon error

lw-lin · zsxwing · commit 569e50680f97 · 2017-01-18T10:52:47.000-08:00
## What changes were proposed in this pull request? We should call `StateStore.abort()` when there should be any error before the store is committed. ## How was this patch tested? Manually. Author: Liwei Lin <lwlin7@gmail.com> Closes #16547 from lw-lin/append-filter.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulAggregate.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulAggregate.scala
@@ -31,6 +31,7 @@ import org.apache.spark.sql.execution.streaming.state._
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.streaming.OutputMode
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.TaskContext
 
 
 /** Used to identify the state store for a given operator. */
@@ -150,6 +151,13 @@ case class StateStoreSaveExec(
         val numTotalStateRows = longMetric("numTotalStateRows")
         val numUpdatedStateRows = longMetric("numUpdatedStateRows")
 
+        // Abort the state store in case of error
+        TaskContext.get().addTaskCompletionListener(_ => {
+          if (!store.hasCommitted) {
+            store.abort()
+          }
+        })
+
         outputMode match {
           // Update and output all rows in the StateStore.
           case Some(Complete) =>
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala
@@ -203,7 +203,7 @@ private[state] class HDFSBackedStateStoreProvider(
     /**
      * Whether all updates have been committed
      */
-    override private[state] def hasCommitted: Boolean = {
+    override private[streaming] def hasCommitted: Boolean = {
       state == COMMITTED
     }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala
@@ -83,7 +83,7 @@ trait StateStore {
   /**
    * Whether all updates have been committed
    */
-  private[state] def hasCommitted: Boolean
+  private[streaming] def hasCommitted: Boolean
 }
 
 

Original file line number	Diff line number	Diff line change
`@@ -203,7 +203,7 @@ private[state] class HDFSBackedStateStoreProvider(`
`203`	`203`	`/**`
`204`	`204`	`* Whether all updates have been committed`
`205`	`205`	`*/`
`206`		`- override private[state] def hasCommitted: Boolean = {`
	`206`	`+ override private[streaming] def hasCommitted: Boolean = {`
`207`	`207`	`state == COMMITTED`
`208`	`208`	`}`
`209`	`209`
Original file line number	Diff line number	Diff line change
`@@ -83,7 +83,7 @@ trait StateStore {`
`83`	`83`	`/**`
`84`	`84`	`* Whether all updates have been committed`
`85`	`85`	`*/`
`86`		`- private[state] def hasCommitted: Boolean`
	`86`	`+ private[streaming] def hasCommitted: Boolean`
`87`	`87`	`}`
`88`	`88`
`89`	`89`