[MINOR][REFACTORING] KeyValueGroupedDataset.mapGroupsWithState uses flatMapGroupsWithState

jaceklaskowski · jaceklaskowski · commit ce5146641198 · 2017-07-15T11:31:55.000+02:00
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala
@@ -242,16 +242,7 @@ class KeyValueGroupedDataset[K, V] private[sql](
   def mapGroupsWithState[S: Encoder, U: Encoder](
       func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
     val flatMapFunc = (key: K, it: Iterator[V], s: GroupState[S]) => Iterator(func(key, it, s))
-    Dataset[U](
-      sparkSession,
-      FlatMapGroupsWithState[K, V, S, U](
-        flatMapFunc.asInstanceOf[(Any, Iterator[Any], LogicalGroupState[Any]) => Iterator[Any]],
-        groupingAttributes,
-        dataAttributes,
-        OutputMode.Update,
-        isMapGroupsWithState = true,
-        GroupStateTimeout.NoTimeout,
-        child = logicalPlan))
+    flatMapGroupsWithState(OutputMode.Update, GroupStateTimeout.NoTimeout)(flatMapFunc)
   }
 
   /**
@@ -278,16 +269,7 @@ class KeyValueGroupedDataset[K, V] private[sql](
       timeoutConf: GroupStateTimeout)(
       func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
     val flatMapFunc = (key: K, it: Iterator[V], s: GroupState[S]) => Iterator(func(key, it, s))
-    Dataset[U](
-      sparkSession,
-      FlatMapGroupsWithState[K, V, S, U](
-        flatMapFunc.asInstanceOf[(Any, Iterator[Any], LogicalGroupState[Any]) => Iterator[Any]],
-        groupingAttributes,
-        dataAttributes,
-        OutputMode.Update,
-        isMapGroupsWithState = true,
-        timeoutConf,
-        child = logicalPlan))
+    flatMapGroupsWithState(OutputMode.Update, timeoutConf)(flatMapFunc)
   }
 
   /**
