[SPARK-36132][SS][SQL] Support initial state for batch mode of flatMapGroupsWithState #33336

rahulsmahadev · 2021-07-14T06:01:29Z

What changes were proposed in this pull request?

Adding support for accepting an initial state with flatMapGroupsWithState in batch mode.

Why are the changes needed?

SPARK-35897 added support for accepting an initial state for streaming queries using flatMapGroupsWithState. the code flow is separate for batch and streaming and required a different PR.

Does this PR introduce any user-facing change?

Yes as discussed above flatMapGroupsWithState in batch mode can accept an initialState, previously this would throw an UnsupportedOperationException

How was this patch tested?

Added relevant unit tests in FlatMapGroupsWithStateSuite and modified the tests JavaDatasetSuite

rahulsmahadev · 2021-07-14T06:01:50Z

cc: @tdas

rahulsmahadev · 2021-07-14T06:02:30Z

sql/core/src/main/scala/org/apache/spark/sql/execution/objects.scala

+
+  override def right: SparkPlan = initialState
+
+  override def outputPartitioning: Partitioning = child.outputPartitioning


This is something im not sure of

what is the output partitioning in MapGroupsExec... it should be the same

its the same

SparkQA · 2021-07-14T07:20:07Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45517/

SparkQA · 2021-07-14T07:55:31Z

Kubernetes integration test status success
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45517/

SparkQA · 2021-07-14T11:02:22Z

Test build #141003 has finished for PR 33336 at commit 71f2302.

This patch passes all tests.
This patch merges cleanly.
This patch adds the following public classes (experimental):
case class MapGroupsWithInitialStateExec(

tdas · 2021-07-14T18:22:57Z

sql/core/src/main/scala/org/apache/spark/sql/execution/objects.scala

+/**
+ * Groups the input rows together and calls the function with each group and an iterator containing
+ * all elements in the group.  The result of this function is flattened before being output. This
+ * version of the Physical operator takes a user provided initial state.


improve the docs as well. "This physical operator is similar to MapGroupsExec but the difference that it calls the user-defined function with .... group + iterator + user-defined initial state"

tdas · 2021-07-14T21:55:20Z

sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala

+          inputData.toDS()
+            .groupByKey(x => x)
+            .flatMapGroupsWithState(
+              Update, GroupStateTimeout.NoTimeout, initialState)(flatMapGroupsWithStateFunc)


shouldnt you be testing this with all timeouts?

tdas · 2021-07-14T21:56:00Z

sql/core/src/main/scala/org/apache/spark/sql/execution/objects.scala

+          val result = func(
+              getKey(keyRow),
+              valueRowIter.map(getValue),
+              groupState)


nit: does it have to be multiline?

tdas · 2021-07-14T21:58:15Z

sql/core/src/main/scala/org/apache/spark/sql/execution/objects.scala

+          var foundInitialStateForKey = false
+          val optionalState = initialStateRowIter.map { initialStateRow =>
+            if (foundInitialStateForKey) {
+              throw new IllegalArgumentException("The initial state provided contained " +


This error should be made into a common error between FlatMapGroupswithStateExec and this class. If possible.

tdas · 2021-07-14T21:59:36Z

sql/core/src/main/scala/org/apache/spark/sql/execution/objects.scala

+        groupedChildDataIter, groupedInitialStateIter, groupingAttributes).flatMap {
+        case (keyRow, valueRowIter, initialStateRowIter) =>
+          var foundInitialStateForKey = false
+          val optionalState = initialStateRowIter.map { initialStateRow =>


add docs here explaining the logic with foundInitialStateForKey

tdas · 2021-07-14T22:09:01Z

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala

+          isFlatMapGroupsWithState, timeout, hasInitialState, initialStateGroupAttrs,
+          initialStateDataAttrs, initialStateDeserializer, initialState, child) =>
+        if (hasInitialState) {
+          execution.MapGroupsWithInitialStateExec(


do you need a separate physical plan? Can you implement this using CoGroupExec

tdas · 2021-07-15T17:24:07Z

...ore/src/main/scala/org/apache/spark/sql/execution/streaming/FlatMapGroupsWithStateExec.scala

+          func, keyDeserializer, valueDeserializer, initialStateDeserializer, groupingAttributes,
+          initialStateGroupAttrs, dataAttributes, initialStateDataAttrs, outputObjAttr,
+          child, initialState) {
+        override def outputPartitioning: Partitioning = child.outputPartitioning


i am actually not sure if this is correct. Even for MapGroupsExec. This operators output rows that are completely new columns... completely different from the key-value columns as input. So the output cannot be partitioned in the same way as the input is.

Its safer to not specify this. If you dont specify, then the worst case is that an extra shuffle will be added if the next operator needs it to be partitioned in some way. If it is specified and it is wrong, then the next operator will end up incorrect assuming that things are partitioned in some way and produce incorrect results.

so you were right about this being fishy.

tdas · 2021-07-15T17:33:04Z

sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala

@@ -1284,6 +1284,9 @@ class FlatMapGroupsWithStateSuite extends StateStoreMetricsTest {
      assertCanGetProcessingTime { state.getCurrentProcessingTimeMs() >= 0 }
      assertCannotGetWatermark { state.getCurrentWatermarkMs() }
      assert(!state.hasTimedOut)
+      if(key.contains("Timeout")) {


nit: space after if

SparkQA · 2021-07-15T18:26:13Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45603/

SparkQA · 2021-07-15T19:02:44Z

Kubernetes integration test status success
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45603/

zsxwing · 2021-07-15T19:02:15Z

...ore/src/main/scala/org/apache/spark/sql/execution/streaming/FlatMapGroupsWithStateExec.scala

+        }.toSeq
+
+        // Create group state object
+        val groupState = GroupStateImpl.createForStreaming(


why do we call createForStreaming rather than createForBatch?

tdas · 2021-07-15T19:58:35Z

sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala

+          inputData.toDS()
+            .groupByKey(x => x)
+            .flatMapGroupsWithState(
+              Update, EventTimeTimeout(), initialState)(flatMapGroupsWithStateFunc)


need testing with all 3 timeout settings.

tdas · 2021-07-15T19:59:26Z

...ore/src/main/scala/org/apache/spark/sql/execution/streaming/FlatMapGroupsWithStateExec.scala

+      val func = (keyRow: Any, values: Iterator[Any], states: Iterator[Any]) => {
+        // Check if there is only one state for every key.
+        var foundInitialStateForKey = false
+        val optionalState = states.map { initialState =>


dont override variables! there is already a initialState function param. its confusing.

tdas · 2021-07-15T21:14:16Z

...ore/src/main/scala/org/apache/spark/sql/execution/streaming/FlatMapGroupsWithStateExec.scala

+  }
+
+  /**
+   * Special handling for when the child relation is a batch relation.


what is the handling here? the doc should say plan logical flatmapGroupsWIthState for batch queries.

tdas · 2021-07-15T21:16:38Z

...ore/src/main/scala/org/apache/spark/sql/execution/streaming/FlatMapGroupsWithStateExec.scala

+   * state is not provided we create an instance of the MapGroupsExec
+   */
+  // scalastyle:off argcount
+  def forBatch(


objectName.forX() generally implies creating instances of type objectName's corresponding class. This is not non-intuitive. rather say "FlatMapGroupsWithStateExec.generateSparkPlanForBatchQueries(...)`

SparkQA · 2021-07-15T22:13:35Z

Test build #141089 has finished for PR 33336 at commit d139950.

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

SparkQA · 2021-07-15T23:30:06Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45615/

SparkQA · 2021-07-16T00:05:43Z

Kubernetes integration test status success
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45615/

SparkQA · 2021-07-16T03:18:25Z

Test build #141104 has finished for PR 33336 at commit 01521ad.

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

tdas · 2021-07-19T23:06:28Z

sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateSuite.scala

+        ("keyOnlyInData"), ("keyInStateAndData-2")
+      )
+      val timeoutMode = timeout match {
+        case "NoTimeout" =>


Why are you doing this?? Why not create seq of objects directly Seq(NoTimeout(), EventTimeTimeout(), ProcessingTimeTimeout()).foreach { timeout =>

If i do that then I need to match on that to get the key right or does timeout have a string method ?

zsxwing · 2021-07-20T22:05:10Z

...ore/src/main/scala/org/apache/spark/sql/execution/streaming/FlatMapGroupsWithStateExec.scala

+        var foundInitialStateForKey = false
+        val optionalState = states.map { stateValue =>
+          if (foundInitialStateForKey) {
+            foundDuplicateInitialKeyException()


This won't catch duplicate keys because states is an Iterator. E.g.,

val states = Iterator(1, 1) var foundInitialStateForKey = false val optionalState = states.map { stateValue => if (foundInitialStateForKey) { throw new RuntimeException("foo") } foundInitialStateForKey = true stateValue }.toSeq optionalState.headOption

the above case won't fail.

You can change toSeq to toArray to fix this issue. Can we add a unit test for this?

oh great catch

does this bug exist for streaming as well?

no we have a test for that flatMapGroupsWithState - initial state - duplicate keys

streaming is using Iterator.foreach.

zsxwing

LGTM pending tests.

SparkQA · 2021-07-20T23:56:29Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45880/

SparkQA · 2021-07-21T00:30:09Z

Kubernetes integration test status success
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45880/

SparkQA · 2021-07-21T01:53:47Z

Kubernetes integration test starting
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45883/

SparkQA · 2021-07-21T02:29:51Z

Kubernetes integration test status success
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45883/

SparkQA · 2021-07-21T04:00:27Z

Test build #141364 has finished for PR 33336 at commit ca77b59.

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

…pGroupsWithState ### What changes were proposed in this pull request? Adding support for accepting an initial state with flatMapGroupsWithState in batch mode. ### Why are the changes needed? SPARK-35897 added support for accepting an initial state for streaming queries using flatMapGroupsWithState. the code flow is separate for batch and streaming and required a different PR. ### Does this PR introduce _any_ user-facing change? Yes as discussed above flatMapGroupsWithState in batch mode can accept an initialState, previously this would throw an UnsupportedOperationException ### How was this patch tested? Added relevant unit tests in FlatMapGroupsWithStateSuite and modified the tests `JavaDatasetSuite` Closes #33336 from rahulsmahadev/flatMapGroupsWithStateBatch. Authored-by: Rahul Mahadev <rahul.mahadev@databricks.com> Signed-off-by: Tathagata Das <tathagata.das1565@gmail.com> (cherry picked from commit efcce23) Signed-off-by: Tathagata Das <tathagata.das1565@gmail.com>

tdas · 2021-07-21T05:53:07Z

merged to master and backported to 3.2 since this is minor change but a good improvement to new feature added in 3.2

SparkQA · 2021-07-21T08:18:43Z

Test build #141367 has finished for PR 33336 at commit 4572809.

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

[SS][SQL] Support initial state for batch mode of flatMapGroupsWithState

71f2302

github-actions bot added SQL STRUCTURED STREAMING labels Jul 14, 2021

rahulsmahadev commented Jul 14, 2021

View reviewed changes

tdas reviewed Jul 14, 2021

View reviewed changes

address commments

d139950

tdas reviewed Jul 15, 2021

View reviewed changes

zsxwing reviewed Jul 15, 2021

View reviewed changes

tdas reviewed Jul 15, 2021

View reviewed changes

address comments

01521ad

tdas reviewed Jul 19, 2021

View reviewed changes

Fix

cb10de4

zsxwing reviewed Jul 20, 2021

View reviewed changes

Fixed

ca77b59

zsxwing approved these changes Jul 20, 2021

View reviewed changes

tdas approved these changes Jul 20, 2021

View reviewed changes

Retrigger tests

4572809

asfgit closed this in efcce23 Jul 21, 2021


		override def right: SparkPlan = initialState

		override def outputPartitioning: Partitioning = child.outputPartitioning

[SPARK-36132][SS][SQL] Support initial state for batch mode of flatMapGroupsWithState #33336

[SPARK-36132][SS][SQL] Support initial state for batch mode of flatMapGroupsWithState #33336

Conversation

rahulsmahadev commented Jul 14, 2021

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

rahulsmahadev commented Jul 14, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SparkQA commented Jul 14, 2021

SparkQA commented Jul 14, 2021

SparkQA commented Jul 14, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SparkQA commented Jul 15, 2021

SparkQA commented Jul 15, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SparkQA commented Jul 15, 2021

SparkQA commented Jul 15, 2021

SparkQA commented Jul 16, 2021

SparkQA commented Jul 16, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zsxwing Jul 20, 2021 • edited Loading

Choose a reason for hiding this comment

zsxwing left a comment

Choose a reason for hiding this comment

SparkQA commented Jul 20, 2021

SparkQA commented Jul 21, 2021

SparkQA commented Jul 21, 2021

SparkQA commented Jul 21, 2021

SparkQA commented Jul 21, 2021

tdas commented Jul 21, 2021

SparkQA commented Jul 21, 2021

zsxwing Jul 20, 2021 •

edited

Loading