[SPARK-35897][SS] Support user defined initial state with flatMapGroupsWithState in Structured Streaming #33093

rahulsmahadev · 2021-06-25T17:42:19Z

What changes were proposed in this pull request?

This PR aims to add support for specifying a user defined initial state for arbitrary structured streaming stateful processing using [flat]MapGroupsWithState operator.

Why are the changes needed?

Users can load previous state of their stateful processing as an initial state instead of redoing the entire processing once again.

Does this PR introduce any user-facing change?

Yes this PR introduces new API

  def mapGroupsWithState[S: Encoder, U: Encoder](
      timeoutConf: GroupStateTimeout,
      initialState: KeyValueGroupedDataset[K, S])(
      func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] 

  def flatMapGroupsWithState[S: Encoder, U: Encoder](
      outputMode: OutputMode,
      timeoutConf: GroupStateTimeout,
      initialState: KeyValueGroupedDataset[K, S])(
      func: (K, Iterator[V], GroupState[S]) => Iterator[U])

How was this patch tested?

Through unit tests in FlatMapGroupsWithStateSuite

rahulsmahadev · 2021-06-25T17:42:55Z

@tdas can you enable CI on this

rahulsmahadev · 2021-06-25T17:45:00Z

cc: @tdas can you enable CI on this

zsxwing · 2021-06-25T18:25:35Z

add to whitelist

SparkQA · 2021-06-25T18:45:32Z

Test build #140337 has finished for PR 33093 at commit 4a0b591.

This patch fails Scala style tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2021-06-25T18:50:36Z

Kubernetes integration test unable to build dist.

exiting with code: 1
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/44868/

SparkQA · 2021-06-25T19:43:24Z

Test build #140342 has finished for PR 33093 at commit 69b3c7e.

This patch fails Scala style tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2021-06-25T19:50:35Z

Kubernetes integration test unable to build dist.

exiting with code: 1
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/44873/

zsxwing · 2021-06-25T21:58:04Z

ok to test

SparkQA · 2021-06-25T22:08:21Z

Kubernetes integration test unable to build dist.

exiting with code: 1
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/44877/

SparkQA · 2021-06-25T22:45:01Z

Test build #140347 has finished for PR 33093 at commit 69b3c7e.

This patch fails Scala style tests.
This patch merges cleanly.
This patch adds no public classes.

HyukjinKwon · 2021-06-28T03:17:57Z

@rahulsmahadev, mind keeping the PR description template (https://github.com/apache/spark/blob/master/.github/PULL_REQUEST_TEMPLATE)?

SparkQA · 2021-06-29T21:36:54Z

Test build #140405 has finished for PR 33093 at commit 6c1443c.

This patch fails Scala style tests.
This patch merges cleanly.
This patch adds the following public classes (experimental):
trait WatermarkSupport

SparkQA · 2021-06-29T22:35:38Z

Test build #140407 has finished for PR 33093 at commit ec2b972.

This patch fails Scala style tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2021-06-30T02:16:46Z

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

SparkQA · 2021-06-30T02:18:09Z

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

SparkQA · 2021-06-30T02:49:17Z

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

SparkQA · 2021-06-30T02:51:48Z

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

SparkQA · 2021-06-30T03:37:15Z

Kubernetes integration test unable to build dist.

exiting with code: 1
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/44928/

SparkQA · 2021-06-30T03:44:22Z

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

SparkQA · 2021-06-30T04:19:29Z

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

SparkQA · 2021-06-30T05:56:42Z

Test build #140413 has finished for PR 33093 at commit f221358.

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

SparkQA · 2021-06-30T06:13:49Z

Test build #140415 has finished for PR 33093 at commit 673008d.

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

SparkQA · 2021-07-01T14:24:18Z

Kubernetes integration test unable to build dist.

exiting with code: 1
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45031/

SparkQA · 2021-07-01T14:30:41Z

Test build #140514 has finished for PR 33093 at commit 57c9c2f.

This patch fails Spark unit tests.
This patch merges cleanly.
This patch adds no public classes.

tdas · 2021-07-01T14:37:04Z

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

+            ("keyInStateAndData-2", Seq[String](), "1"),
+            ("keyOnlyInData", Seq[String]("keyOnlyInData"), "1") // inc by 1
+          ),
+          assertNumStateRows(total = 5, updated = 5),


You are not testing whether the initial group state is actually being saved or not. you could have just created th e input GroupState object with the initial state and not saved to state store, and this test will still pass. So you need to run another batch to retrieve and test the save state.

Furthermore, you need to explicitly test whether the initial state is saved to store even if you dont call GroupState.update(). Right now in your test function, you are always calling update. So even if you incorrectly did not save the initial state store, the update will always make sure the state store is updated. So you need to test for more cases, with more keys.

SparkQA · 2021-07-01T16:19:45Z

Test build #140518 has finished for PR 33093 at commit 4e85a62.

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

SparkQA · 2021-07-01T21:31:12Z

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

SparkQA · 2021-07-01T22:04:17Z

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

rahulsmahadev · 2021-07-01T22:52:16Z

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

  def testWithTimeout(timeoutConf: GroupStateTimeout): Unit = {
    test("SPARK-20714: watermark does not fail query when timeout = " + timeoutConf) {
      // 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, Long)], state: GroupState[RunningCount]) => {
+          (key: String, values: Iterator[(String, Long)], state: GroupState[RunningCount]) => {


nit: fix this

SparkQA · 2021-07-01T23:29:27Z

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

tdas · 2021-07-01T23:49:31Z

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

+        assert(state.exists)
+        assertCanGetProcessingTime { state.getCurrentProcessingTimeMs() >= 0 }
+        assertCannotGetWatermark { state.getCurrentWatermarkMs() }
+        assert(!state.hasTimedOut)


why does the last 3 tests need to be inside if (valList.isEmpty) {?

tdas · 2021-07-01T23:53:31Z

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

+      // We need to check if not explicitly calling update will still save the init state or not
+      if (valList.nonEmpty || state.getOption.map(_.count).getOrElse(0L) != 2L) {
+        // this is not reached when valList is empty and the state count is 2
+        state.update(new RunningCount(count))


rather than doing this complicated logic of no updating when some magical set of conditions are met .. isnt it simpler to have

if (!key.contains("NoUpdate")) state.update(...)

and then pass a key name keyOnlyInStateButNoUpate or keyInStateAndDataButNoUpate??

tdas · 2021-07-01T23:55:44Z

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

+            ("keyOnlyInState-1", Seq[String](), "1"),
+            ("keyOnlyInState-2", Seq[String](), "2"),
+            ("keyInStateAndData-2", Seq[String]("keyInStateAndData-2"), "3"), // inc by 1
+            ("keyInStateAndData-1", Seq[String](), "1"),


keyInStateAndData-1 is NOT in data. this is confusing!

hmm it is added in second batch

this is nit, the point naming it like this is its in both state and data of the first batch which is what we are mainly testing. hence it is confusing.

tdas · 2021-07-01T23:56:27Z

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

+        assert(state.exists)
+        assertCanGetProcessingTime { state.getCurrentProcessingTimeMs() >= 0 }
+        assertCannotGetWatermark { state.getCurrentWatermarkMs() }
+        assert(!state.hasTimedOut)


same questions as above

SparkQA · 2021-07-01T23:59:41Z

Test build #140545 has finished for PR 33093 at commit b47ac23.

This patch fails Java style tests.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2021-07-02T00:05:02Z

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

SparkQA · 2021-07-02T00:21:45Z

Kubernetes integration test unable to build dist.

exiting with code: 1
URL: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder-K8s/45058/

SparkQA · 2021-07-02T00:59:23Z

Test build #140538 has finished for PR 33093 at commit 0374bdb.

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

SparkQA · 2021-07-02T01:27:34Z

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

SparkQA · 2021-07-02T02:01:17Z

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

SparkQA · 2021-07-02T03:26:53Z

Test build #140543 has finished for PR 33093 at commit b8c70ab.

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

gengliangwang · 2021-07-02T03:52:46Z

Thanks, merging to master

SparkQA · 2021-07-02T05:01:04Z

Test build #140550 has finished for PR 33093 at commit eb83b68.

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

github-actions bot added SQL STRUCTURED STREAMING labels Jun 25, 2021

HyukjinKwon marked this pull request as draft June 28, 2021 03:17

FlatMapGroupsWith state accept initial state

6c1443c

rahulsmahadev force-pushed the flatMapGroupsWithState branch from 69b3c7e to 6c1443c Compare June 29, 2021 20:45

Fix build

ec2b972

rahulsmahadev added 2 commits June 29, 2021 18:30

makes tests cleaner

f221358

simplify test

673008d

tdas reviewed Jul 1, 2021

View reviewed changes

zsxwing marked this pull request as ready for review July 1, 2021 19:57

address comments

0374bdb

aadd java tests and remove extra tests

b8c70ab

rahulsmahadev commented Jul 1, 2021

View reviewed changes

Add test

b47ac23

tdas reviewed Jul 1, 2021

View reviewed changes

Address comments

eb83b68

tdas approved these changes Jul 2, 2021

View reviewed changes

gengliangwang approved these changes Jul 2, 2021

View reviewed changes

gengliangwang closed this in 47485a3 Jul 2, 2021

[SPARK-35897][SS] Support user defined initial state with flatMapGroupsWithState in Structured Streaming #33093

[SPARK-35897][SS] Support user defined initial state with flatMapGroupsWithState in Structured Streaming #33093

Conversation

rahulsmahadev commented Jun 25, 2021 • edited Loading

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 Jun 25, 2021

rahulsmahadev commented Jun 25, 2021

zsxwing commented Jun 25, 2021

SparkQA commented Jun 25, 2021

SparkQA commented Jun 25, 2021

SparkQA commented Jun 25, 2021

SparkQA commented Jun 25, 2021

zsxwing commented Jun 25, 2021

SparkQA commented Jun 25, 2021

SparkQA commented Jun 25, 2021

HyukjinKwon commented Jun 28, 2021

SparkQA commented Jun 29, 2021

SparkQA commented Jun 29, 2021

SparkQA commented Jun 30, 2021

SparkQA commented Jun 30, 2021

SparkQA commented Jun 30, 2021

SparkQA commented Jun 30, 2021

SparkQA commented Jun 30, 2021

SparkQA commented Jun 30, 2021

SparkQA commented Jun 30, 2021

SparkQA commented Jun 30, 2021

SparkQA commented Jun 30, 2021

SparkQA commented Jul 1, 2021

SparkQA commented Jul 1, 2021

tdas Jul 1, 2021 • edited Loading

Choose a reason for hiding this comment

SparkQA commented Jul 1, 2021

SparkQA commented Jul 1, 2021

SparkQA commented Jul 1, 2021

rahulsmahadev Jul 1, 2021

Choose a reason for hiding this comment

SparkQA commented Jul 1, 2021

tdas Jul 1, 2021

Choose a reason for hiding this comment

tdas Jul 1, 2021

Choose a reason for hiding this comment

tdas Jul 1, 2021

Choose a reason for hiding this comment

rahulsmahadev Jul 2, 2021

Choose a reason for hiding this comment

tdas Jul 2, 2021

Choose a reason for hiding this comment

tdas Jul 1, 2021

Choose a reason for hiding this comment

SparkQA commented Jul 1, 2021

SparkQA commented Jul 2, 2021

SparkQA commented Jul 2, 2021

SparkQA commented Jul 2, 2021

SparkQA commented Jul 2, 2021

SparkQA commented Jul 2, 2021

SparkQA commented Jul 2, 2021

gengliangwang commented Jul 2, 2021

SparkQA commented Jul 2, 2021

rahulsmahadev commented Jun 25, 2021 •

edited

Loading

tdas Jul 1, 2021 •

edited

Loading