[SPARK-19413][SS] MapGroupsWithState for arbitrary stateful operations #16758

tdas · 2017-01-31T19:16:30Z

What changes were proposed in this pull request?

mapGroupsWithState is a new API for arbitrary stateful operations in Structured Streaming, similar to DStream.mapWithState

Requirements

Users should be able to specify a function that can do the following
Access the input row corresponding to a key
Access the previous state corresponding to a key
Optionally, update or remove the state
Output any number of new rows (or none at all)

Proposed API

// ------------ New methods on KeyValueGroupedDataset ------------
class KeyValueGroupedDataset[K, V] {	
	// Scala friendly
	def mapGroupsWithState[S: Encoder, U: Encoder](func: (K, Iterator[V], KeyedState[S]) => U)
        def flatMapGroupsWithState[S: Encode, U: Encoder](func: (K, Iterator[V], KeyedState[S]) => Iterator[U])
	// Java friendly
       def mapGroupsWithState[S, U](func: MapGroupsWithStateFunction[K, V, S, R], stateEncoder: Encoder[S], resultEncoder: Encoder[U])
       def flatMapGroupsWithState[S, U](func: FlatMapGroupsWithStateFunction[K, V, S, R], stateEncoder: Encoder[S], resultEncoder: Encoder[U])
}

// ------------------- New Java-friendly function classes ------------------- 
public interface MapGroupsWithStateFunction<K, V, S, R> extends Serializable {
  R call(K key, Iterator<V> values, state: KeyedState<S>) throws Exception;
}
public interface FlatMapGroupsWithStateFunction<K, V, S, R> extends Serializable {
  Iterator<R> call(K key, Iterator<V> values, state: KeyedState<S>) throws Exception;
}

// ---------------------- Wrapper class for state data ---------------------- 
trait KeyedState[S] {
	def exists(): Boolean  	
  	def get(): S 			// throws Exception is state does not exist
	def getOption(): Option[S]       
	def update(newState: S): Unit
	def remove(): Unit		// exists() will be false after this
}

Key Semantics of the State class

The state can be null.
If the state.remove() is called, then state.exists() will return false, and getOption will returm None.
After that state.update(newState) is called, then state.exists() will return true, and getOption will return Some(...).
None of the operations are thread-safe. This is to avoid memory barriers.

Usage

val stateFunc = (word: String, words: Iterator[String, runningCount: KeyedState[Long]) => {
    val newCount = words.size + runningCount.getOption.getOrElse(0L)
    runningCount.update(newCount)
   (word, newCount)
}

dataset					                        // type is Dataset[String]
  .groupByKey[String](w => w)        	                // generates KeyValueGroupedDataset[String, String]
  .mapGroupsWithState[Long, (String, Long)](stateFunc)	// returns Dataset[(String, Long)]

How was this patch tested?

New unit tests.

SparkQA · 2017-01-31T19:22:20Z

Test build #72205 has finished for PR 16758 at commit 6fab7a5.

This patch fails to build.
This patch merges cleanly.
This patch adds no public classes.

tdas · 2017-01-31T19:22:57Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulAggregate.scala

  }
 }

+trait StateStoreReader extends StatefulOperator {


This files should probably be renamed from StatefulAggregation to StatefulOperations.

and lowercase if it contains multiple classes

SparkQA · 2017-01-31T19:58:18Z

Test build #72206 has finished for PR 16758 at commit 8be63de.

This patch fails to build.
This patch merges cleanly.
This patch adds no public classes.

marmbrus

Looking pretty good!

marmbrus · 2017-01-31T21:25:53Z

sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala

+   * optionally update or remove the corresponding state. The returned object will form a new
+   * [[Dataset]].
+   *
+   * This function can be applied on both batch and streaming Datasets. With a streaming dataset,


will be called

marmbrus · 2017-01-31T21:30:28Z

sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala

+   * [[Dataset]].
+   *
+   * This function can be applied on both batch and streaming Datasets. With a streaming dataset,
+   * this function will be once for each in every trigger. For each key, the updated state from the


Any updates to the state will be stored and passed to the user given function in subsequent batches when executed as a Streaming Query.

marmbrus · 2017-01-31T21:32:59Z

sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala


+  /**
+   * ::Experimental::
+   * (Scala-specific)


while maintaining some user-defined state for each key.

marmbrus · 2017-01-31T21:35:08Z

sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala

+  /**
+   * ::Experimental::
+   * (Scala-specific)
+   * Applies the given function to each group of data, while using an additional keyed state.


I would consider breaking this up to make it a little easier to follow:

For each unique group, the given function will be invoked with the following arguments: - The key of the group. - A user-defined state object set by previous invocations of the given function. Note that, for batch queries, there is only ever one invocation and thus the state object will always be empty. - An iterator containing all the values for this key.

marmbrus · 2017-01-31T21:36:07Z

sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala

+   * function call in a trigger will be the state available in the function call in the next
+   * trigger. However, for batch, `mapGroupsWithState` behaves exactly as `mapGroups` and the
+   * function is called only once per key without any prior state.
+   *


I'd maybe put these into bullets as well.

marmbrus · 2017-01-31T21:49:55Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulAggregate.scala

+    ClusteredDistribution(groupingAttributes) :: Nil
+
+  override def requiredChildOrdering: Seq[Seq[SortOrder]] =
+    Seq(groupingAttributes.map(SortOrder(_, Ascending)))   // is this ordering needed?


Yes, the GroupedIterator relies on sorting.

Thanks for confirming, will remove the comment.

marmbrus · 2017-01-31T21:50:19Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulAggregate.scala

+    Seq(groupingAttributes.map(SortOrder(_, Ascending)))   // is this ordering needed?
+
+  override protected def doExecute(): RDD[InternalRow] = {
+


nit: extra newline

marmbrus · 2017-01-31T21:51:10Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulAggregate.scala

+      child.output.toStructType,
+      sqlContext.sessionState,
+      Some(sqlContext.streams.stateStoreCoordinator)) { (store, iter) =>
+        try {


existing, but should mapPartitionsWithStateStore be implementing the abort handling. Seems generic.

I thought of that. If I change the mapPartitionsWithStateStore, then StateStoreSave should also need to be change, which I didnt want to touch in this already big PR

Managed to do it. Not a big change. Improves correctness, and reduces code.

marmbrus · 2017-01-31T21:52:54Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulAggregate.scala


            // Assumption: Append mode can be done only when watermark has been specified
-            store.remove(watermarkPredicate.get.eval)
+            store.remove(watermarkPredicate.get.eval _)


Why this change?

Because of the addition of a new version of stateStore.remove. I think the compiler can disambiguate correctly without this.

marmbrus · 2017-01-31T21:54:01Z

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

  }

+  def deserializeRowToObject(
+    deserializer: Expression): InternalRow => Any = {


indent, also does this not fit?

SparkQA · 2017-01-31T23:42:21Z

Test build #72215 has finished for PR 16758 at commit 34449e4.

This patch fails to build.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2017-01-31T23:48:43Z

Test build #72216 has finished for PR 16758 at commit 3628af8.

This patch fails to build.
This patch merges cleanly.
This patch adds no public classes.

SparkQA · 2017-02-01T02:06:48Z

Test build #72217 has finished for PR 16758 at commit 59c229b.

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

SparkQA · 2017-02-01T02:53:46Z

Test build #72223 has finished for PR 16758 at commit 3dc0353.

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

SparkQA · 2017-02-01T06:56:39Z

Test build #72230 has finished for PR 16758 at commit 8b18fa1.

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

SparkQA · 2017-02-01T22:05:45Z

Test build #72256 has finished for PR 16758 at commit 8b3150a.

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

lw-lin · 2017-02-02T06:37:48Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/KeyedStateImpl.scala

+    removed = true
+  }
+
+  override def toString: String = "KeyedState($value)"


nit: s"KeyedState($value)"
:)

lw-lin · 2017-02-02T07:10:48Z

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

+    // Function that returns running count only if its even, otherwise does not return
+    val stateFunc = (key: String, values: Iterator[String], state: KeyedState[RunningCount]) => {
+      if (state.exists) throw new IllegalArgumentException("state.exists should be false")
+      if (state.exists) {


nit: state.get == ...? also state.get in the error message

done. not a nit! wrong test!

lw-lin · 2017-02-02T07:14:58Z

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

+  test("mapGroupsWithState - batch") {
+    val stateFunc = (key: String, values: Iterator[String], state: KeyedState[RunningCount]) => {
+      if (state.exists) throw new IllegalArgumentException("state.exists should be false")
+      if (state.exists) {


done. thanks for catching these.

lw-lin · 2017-02-02T07:19:46Z

sql/core/src/main/scala/org/apache/spark/sql/KeyedState.scala

+ *  - If the `remove()` is called, then `exists()` will return `false`, and
+ *    `getOption()` will return `None`.
+ *  - After that `update(newState)` is called, then `exists()` will return `true`,
+ *    and `getOption()` will return `Some(...)`.


nit: getOption ...

zsxwing

Made one pass. Most of my comments are nits.

zsxwing · 2017-02-02T21:33:01Z

...alyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationsSuite.scala


+  // MapGroupsWithState: Not supported after a streaming aggregation
+  val att = new AttributeReference(name = "a", dataType = LongType)()
+  assertSupportedInStreamingPlan(


nit: assertSupportedInStreamingPlan -> assertSupportedInBatchPlan

zsxwing · 2017-02-02T21:50:50Z

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

+    // Returns the data and the count if state is defined, otherwise does not return anything
+    val stateFunc = (key: String, values: Iterator[String], state: KeyedState[RunningCount]) => {
+
+      var count = Option(state.get).map(_.count).getOrElse(0L) + values.size


nit: var -> val

zsxwing · 2017-02-02T22:05:42Z

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

+      StartStream(),
+      CheckLastBatch(("a", 3L))     // task should not fail, and should show correct count
+    )
+  }


Could you add a test for aggregation after mapGroupsWithState?

zsxwing · 2017-02-02T22:36:42Z

sql/core/src/main/scala/org/apache/spark/sql/KeyedState.scala

+
+  /**
+   * Update the value of the state. Note that null is not a valid value, and `update(null)` is
+   * same as `remove()`


It's better to disallow this case. Otherwise, the user may happen to send a null by mistake and we just hide the error.

@marmbrus any thoughts on this?

if get for a non-existent key returns null then I think its reasonable that remove(key) and put(key, null) are the same.

@marmbrus The user may pass a null by mistake (e.g., just call some method that may return null but not be aware of it). It's pretty hard to debug such silent mistake.

zsxwing · 2017-02-02T22:38:10Z

sql/core/src/main/scala/org/apache/spark/sql/KeyedState.scala

+ *  - The key of the group.
+ *  - An iterator containing all the values for this key.
+ *  - A user-defined state object set by previous invocations of the given function.
+ * In case of a batch Dataset, there is only invocation and state object will be empty as


nit: only one invocation

zsxwing · 2017-02-02T22:48:28Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala

   * preparation walks the query plan.
   */
-  private var operatorId = 0
+  private val operatorId = new AtomicInteger(0)


there is only one thread here. Why change it? Any concern?

not really. just a tiny bit cleaner to do getAndIncrement than using x and then calling x += 1

zsxwing · 2017-02-02T22:54:13Z

...main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala

+          case Some(ValueRemoved(_, _)) =>
+          // Remove already in update map, no need to change
+        }
+      }


Missing writeToDeltaFile(tempDeltaFileStream, ValueRemoved(key, value)). It's better to extract the duplicated codes to a new method.

yeah. that's true. I should add tests for this StateStore operation.

zsxwing · 2017-02-02T22:54:28Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

   */
  def remove(condition: UnsafeRow => Boolean): Unit

+  def remove(key: UnsafeRow): Unit


nit: missing scala doc

zsxwing · 2017-02-02T23:10:45Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala

+        val groupedIter = GroupedIterator(iter, groupingAttributes, child.output)
+
+        val getKeyObj = ObjectOperator.deserializeRowToObject(keyDeserializer, groupingAttributes)
+        val getKey = GenerateUnsafeProjection.generate(groupingAttributes, child.output)


nit: not used

zsxwing · 2017-02-02T23:13:37Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala

+          val wrappedState = KeyedStateImpl[Any](stateObjOption.orNull)
+          val mappedIterator = func(keyObj, valueObjIter, wrappedState)
+
+          if (wrappedState.isRemoved) {


this should not be checked here. mappedIterator may be lazy. You can wrap it with a CompletionIterator and add these codes there.

good point. this would confusing for the user if we dont do this.

brkyvz · 2017-02-03T00:36:10Z

...lyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnsupportedOperationChecker.scala


+    /** Collect all the streaming aggregates in a sub plan */
+    def collectStreamingAggregates(subplan: LogicalPlan): Seq[Aggregate] = {
+      subplan.collect { case a@Aggregate(_, _, _) if a.isStreaming => a }


nit: space before and after @. Actually if you're not going to use any of the Aggregate params, just change this to a: Aggregate

brkyvz · 2017-02-03T00:37:30Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/object.scala

+    stateSerializer: Seq[NamedExpression],
+    child: LogicalPlan) extends UnaryNode with ObjectProducer
+
+


nit: extra line

brkyvz · 2017-02-03T00:47:33Z

sql/core/src/main/scala/org/apache/spark/sql/KeyedState.scala

+ * Important points to note about using KeyedState.
+ *  - The value of the state cannot be null. So updating state with null is same as removing it.
+ *  - Operations on `KeyedState` are not thread-safe. This is to avoid memory barriers.
+ *  - If the `remove()` is called, then `exists()` will return `false`, and


nit: If ``remove`` is called. Remove the

brkyvz · 2017-02-03T00:54:42Z

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/KeyedStateImpl.scala

+
+import org.apache.spark.sql.KeyedState
+
+/** Internal implementation of the [[KeyedState]] interface */


Would be nice to note here that this implementation is not thread safe

i have mentioned that the trait KeyedState is not thread-safe

tdas · 2017-02-07T06:23:40Z

I addressed all the comments. However, @zsxwing @marmbrus, our offline discussion of throwing error on .update(null) ran into a problem. Since its typed as S, the behavior is odd when S is primitive type. See the failing test. When the type is Int, get return 0 when the state does not exist. That's very non-intuitive.

SparkQA · 2017-02-07T06:23:42Z

Test build #72490 has started for PR 16758 at commit b32dcd1.

tdas · 2017-02-07T22:20:23Z

@zsxwing had a discussion with @marmbrus. It is indeed weird that after calling state.remove() and state.get will return 0 if it is of type int (as in JVM null casted to int = 0). So we went back to the semantics of get throwing error when state does not exist.
This ensures the following semantics that @marmbrus originally desired.

update(null) is not allowed - for primitive types, compile disallows it; for other types, it throw IllegalArgumentException
Hence, get will never return Some(null).

SparkQA · 2017-02-08T00:56:43Z

Test build #72543 has finished for PR 16758 at commit f3d1231.

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

zsxwing · 2017-02-08T04:20:12Z

LGTM. Merging to master and ~~2.1~~.

zsxwing · 2017-02-08T04:24:01Z

It conflicts with 2.1. Could you submit a backport PR, please?

tdas added 12 commits January 17, 2017 16:49

Prototype - almost working

d10da27

Renamed to mapGroupsWithState

78cd185

Fixed bugs

0c22e08

Removed prints

52e14e4

Test state remove

529aefe

Test restart, and test with metrics

57f5e8d

Fixed everything

3e0d8dc

Refactored, added java APIs and tests

b54fa23

Refactored

ab3cb6c

Added more test

ddf4550

Merge remote-tracking branch 'apache-github/master' into mapWithState

3133f83

Added docs

6fab7a5

tdas changed the title ~~Arbitrary stateful operations with MapGroupsWithState~~ [SPARK-19413][SS] MapGroupsWithState for arbitrary stateful operations Jan 31, 2017

tdas commented Jan 31, 2017

View reviewed changes

Added docs

8be63de

marmbrus reviewed Jan 31, 2017

View reviewed changes

tdas added 2 commits January 31, 2017 15:33

Addressed many comments

34449e4

Addressed more comments

3628af8

tdas added 4 commits January 31, 2017 16:02

Addresed more comments

af4f1f2

Added missing classes

59c229b

Addressed 1 comment

db2dbb2

Merge remote-tracking branch 'apache-github/master' into mapWithState

3dc0353

Addressed comments

8b18fa1

Fixed bug

8b3150a

lw-lin reviewed Feb 2, 2017

View reviewed changes

zsxwing requested changes Feb 2, 2017

View reviewed changes

brkyvz reviewed Feb 3, 2017

View reviewed changes

tdas added 2 commits February 7, 2017 01:17

Addressed concerns

7a39eaf

Merge remote-tracking branch 'apache-github/master' into mapWithState

b32dcd1

Addressed comments

f3d1231

asfgit closed this in aeb8034 Feb 8, 2017

tdas mentioned this pull request Feb 8, 2017

[SPARK-19413][SS] MapGroupsWithState for arbitrary stateful operations for branch-2.1 #16850

Closed

		Seq(groupingAttributes.map(SortOrder(_, Ascending))) // is this ordering needed?

		override protected def doExecute(): RDD[InternalRow] = {

		stateSerializer: Seq[NamedExpression],
		child: LogicalPlan) extends UnaryNode with ObjectProducer


		import org.apache.spark.sql.KeyedState

		/** Internal implementation of the [[KeyedState]] interface */

[SPARK-19413][SS] MapGroupsWithState for arbitrary stateful operations #16758

[SPARK-19413][SS] MapGroupsWithState for arbitrary stateful operations #16758

Uh oh!

Conversation

tdas commented Jan 31, 2017 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

How was this patch tested?

Uh oh!

SparkQA commented Jan 31, 2017

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SparkQA commented Jan 31, 2017

Uh oh!

marmbrus left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SparkQA commented Jan 31, 2017

Uh oh!

SparkQA commented Jan 31, 2017

Uh oh!

SparkQA commented Feb 1, 2017

Uh oh!

SparkQA commented Feb 1, 2017

Uh oh!

SparkQA commented Feb 1, 2017

Uh oh!

SparkQA commented Feb 1, 2017

Uh oh!

lw-lin Feb 2, 2017 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lw-lin Feb 2, 2017 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zsxwing left a comment

Choose a reason for hiding this comment

tdas commented Jan 31, 2017 •

edited

Loading

lw-lin Feb 2, 2017 •

edited

Loading

lw-lin Feb 2, 2017 •

edited

Loading