tdas commented on a change in pull request #27333: [SPARK-29438][SS][FOLLOWUP]
Add regression tests for Streaming Aggregation and flatMapGroupsWithState
URL: https://github.com/apache/spark/pull/27333#discussion_r373322828
##########
File path:
sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala
##########
@@ -1019,6 +1019,76 @@ class FlatMapGroupsWithStateSuite extends
StateStoreMetricsTest {
spark.createDataset(Seq(("a", 2), ("b", 1))).toDF)
}
+ testWithAllStateVersions("SPARK-29438: ensure UNION doesn't lead
(flat)MapGroupsWithState" +
+ " to use shifted partition IDs") {
+ // Function to maintain running count up to 2, and then remove the count
+ // Returns the data and the count (-1 if count reached beyond 2 and state
was just removed)
+ val stateFunc = (key: String, values: Iterator[String], state:
GroupState[RunningCount]) => {
+ assertCanGetProcessingTime { state.getCurrentProcessingTimeMs() >= 0 }
+ assertCannotGetWatermark { state.getCurrentWatermarkMs() }
+
+ val count = state.getOption.map(_.count).getOrElse(0L) + values.size
+ if (count == 3) {
+ state.remove()
+ (key, "-1")
+ } else {
+ state.update(RunningCount(count))
+ (key, count.toString)
+ }
+ }
+
+ def constructUnionDf(desiredPartitionsForInput1: Int)
+ : (MemoryStream[String], MemoryStream[String], DataFrame) = {
+ val input1 = MemoryStream[String](desiredPartitionsForInput1)
+ val input2 = MemoryStream[String]
+ val df1 = input1.toDF()
+ .select($"value", $"value")
+ val df2 = input2.toDS()
+ .groupByKey(x => x)
+ .mapGroupsWithState(stateFunc) // Types = State: MyState, Out: (Str,
Str)
+ .toDF()
+
+ // Unioned DF would have columns as (String, String)
+ (input1, input2, df1.union(df2))
+ }
+
+ withTempDir { checkpointDir =>
+ val (input1, input2, unionDf) = constructUnionDf(2)
+ testStream(unionDf, Update)(
+ StartStream(checkpointLocation = checkpointDir.getAbsolutePath),
+ MultiAddData(
+ (input1, Seq("a")),
+ (input2, Seq("a"))),
+ CheckNewAnswer(("a", "a"), ("a", "1")),
Review comment:
input1 is just unioned at the end and does not need to share the same keys
as input2, i would make the input of input1 very different from input2, such
that it is obvious in the CheckAnswer which output came from where. For example
input1's strings could "input1-a" and input2's can be "input2-a"
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
