[SPARK-23004][SS] Ensure StateStore.commit is called only once in a streaming aggregation task

[tdas]

What changes were proposed in this pull request?

A structured streaming query with a streaming aggregation can throw the following error in rare cases. 

java.lang.IllegalStateException: Cannot commit after already committed or aborted
at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider.org$apache$spark$sql$execution$streaming$state$HDFSBackedStateStoreProvider$$verify(HDFSBackedStateStoreProvider.scala:643)
at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider$HDFSBackedStateStore.commit(HDFSBackedStateStoreProvider.scala:135)
at org.apache.spark.sql.execution.streaming.StateStoreSaveExec$$anonfun$doExecute$3$$anon$2$$anonfun$hasNext$2.apply$mcV$sp(statefulOperators.scala:359)
at org.apache.spark.sql.execution.streaming.StateStoreWriter$class.timeTakenMs(statefulOperators.scala:102)
at org.apache.spark.sql.execution.streaming.StateStoreSaveExec.timeTakenMs(statefulOperators.scala:251)
at org.apache.spark.sql.execution.streaming.StateStoreSaveExec$$anonfun$doExecute$3$$anon$2.hasNext(statefulOperators.scala:359)
at org.apache.spark.sql.execution.aggregate.ObjectAggregationIterator.processInputs(ObjectAggregationIterator.scala:188)
at org.apache.spark.sql.execution.aggregate.ObjectAggregationIterator.<init>(ObjectAggregationIterator.scala:78)
at org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec$$anonfun$doExecute$1$$anonfun$2.apply(ObjectHashAggregateExec.scala:114)
at org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec$$anonfun$doExecute$1$$anonfun$2.apply(ObjectHashAggregateExec.scala:105)
at org.apache.spark.rdd.RDD$$anonfun$mapPartitionsWithIndexInternal$1$$anonfun$apply$24.apply(RDD.scala:830)
at org.apache.spark.rdd.RDD$$anonfun$mapPartitionsWithIndexInternal$1$$anonfun$apply$24.apply(RDD.scala:830)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:42)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:336)

This can happen when the following conditions are accidentally hit. 

• Streaming aggregation with aggregation function that is a subset of TypedImperativeAggregation (for example, collect_set, collect_list, percentile, etc.). 
• Query running in update} mode
• After the shuffle, a partition has exactly 128 records. 

This causes StateStore.commit to be called twice. See the JIRA for a more detailed explanation. The solution is to use NextIterator or CompletionIterator, each of which has a flag to prevent the "onCompletion" task from being called more than once. In this PR, I chose to implement using NextIterator.

How was this patch tested?

Added unit test that I have confirm will fail without the fix.

[tdas]

@brkyvz PTAL.

[tedyu]

+1

[ConcurrencyPractitioner]

+1

[SparkQA]

Test build #89694 has finished for PR 21124 at commit 304498e.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[brkyvz]

LGTM!