You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "bogao007 (via GitHub)" <gi...@apache.org> on 2023/06/12 18:53:16 UTC
[GitHub] [spark] bogao007 opened a new pull request, #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
bogao007 opened a new pull request, #41558:
URL: https://github.com/apache/spark/pull/41558
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'core/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
To support streaming state APIs in Spark Connect
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as the documentation fix.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
yes
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
Added unit test
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227293576
##########
connector/connect/common/src/main/java/org/apache/spark/api/java/function/FlatMapGroupsWithStateFunction.java:
##########
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.streaming.GroupState;
+
+/**
+ * ::Experimental::
+ * Base interface for a map function used in
+ * {@code org.apache.spark.sql.KeyValueGroupedDataset.flatMapGroupsWithState(
+ * FlatMapGroupsWithStateFunction, org.apache.spark.sql.streaming.OutputMode,
+ * org.apache.spark.sql.Encoder, org.apache.spark.sql.Encoder)}
+ * @since 2.1.1
Review Comment:
nit: version
##########
connector/connect/common/src/main/java/org/apache/spark/api/java/function/MapGroupsWithStateFunction.java:
##########
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.streaming.GroupState;
+
+/**
+ * ::Experimental::
+ * Base interface for a map function used in
+ * {@code org.apache.spark.sql.KeyValueGroupedDataset.mapGroupsWithState(
+ * MapGroupsWithStateFunction, org.apache.spark.sql.Encoder, org.apache.spark.sql.Encoder)}
+ * @since 2.1.1
Review Comment:
nit: version
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -572,6 +726,101 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV](
agg(aggregator)
}
+ override def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
Review Comment:
nit: indentation
##########
connector/connect/common/src/main/scala/org/apache/spark/sql/connect/common/UdfUtils.scala:
##########
@@ -95,6 +96,26 @@ private[sql] object UdfUtils extends Serializable {
}
}
+ def mapValuesAdaptor[K, V, S, U, IV](
+ f: (K, Iterator[V], GroupState[S]) => U,
+ valueMapFunc: IV => V): (K, Iterator[IV], GroupState[S]) => U = {
+ (k: K, itr: Iterator[IV], s: GroupState[S]) =>
+ {
+ f(k, itr.map(v => valueMapFunc(v)), s)
+ }
+ }
+
+ def mapGroupsWithStateFuncToScalaFunc[K, V, S, U](
+ f: MapGroupsWithStateFunction[K, V, S, U]): (K, Iterator[V], GroupState[S]) => U = {
+ (key, data, groupState) => f.call(key, data.asJava, groupState)
+ }
+
+ def flatMapGroupsWithStateFuncToScalaFunc[K, V, S, U](
Review Comment:
With this function, I think we also need to add `java-specific` APIs to `KeyValueGroupedDataset`?
For example:
https://github.com/apache/spark/blob/28869401eff8fcd94aa68d662fdd4d3dbe2f58af/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala#L118-L136
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1230203950
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -460,6 +461,159 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
cogroupSorted(other)(thisSortExprs: _*)(otherSortExprs: _*)(
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ mapGroupsWithState(GroupStateTimeout.NoTimeout)(func)
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param timeoutConf
+ * Timeout configuration for groups that do not receive data for a while.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
Review Comment:
Updated
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1231423053
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -462,6 +462,15 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+ protected def flatMapGroupsWithStateHelper[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S],
Review Comment:
I see. I'm good with keeping the current implementation. Thanks for the explanation!
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1231370155
##########
connector/connect/common/src/main/protobuf/spark/connect/relations.proto:
##########
@@ -858,6 +857,21 @@ message GroupMap {
// (Optional) Expressions for sorting. Only used by Scala Sorted Group Map API.
repeated Expression sorting_expressions = 4;
+
+ // (Optional) Input relation for initial State.
+ Relation initial_input = 5;
Review Comment:
Seems these fields are only used in (flat)mapGroupsWithState, could we add comment above?
##########
connector/connect/common/src/main/scala/org/apache/spark/sql/connect/common/UdfUtils.scala:
##########
@@ -97,22 +97,27 @@ private[sql] object UdfUtils extends Serializable {
}
def mapValuesAdaptor[K, V, S, U, IV](
- f: (K, Iterator[V], GroupState[S]) => U,
- valueMapFunc: IV => V): (K, Iterator[IV], GroupState[S]) => U = {
- (k: K, itr: Iterator[IV], s: GroupState[S]) =>
- {
+ f: (K, Iterator[V], GroupState[S]) => Iterator[U],
+ valueMapFunc: IV => V): (K, Iterator[IV], GroupState[S]) => Iterator[U] = {
+ print("Used function: mapValuesAdaptor")
Review Comment:
I think we should delete this
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -462,6 +462,15 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+ protected def flatMapGroupsWithStateHelper[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S],
Review Comment:
I feel this helper method shouldn't appear there. Is this change because of this comment? https://github.com/apache/spark/pull/41558#discussion_r1227337649
I actually think it's good to use the original implementation, i.e. have a dedicated `flatMapGroupsWithState` and `mapGroupsWIthState` that handles all cases respectively, this also follows the original API design: https://github.com/apache/spark/blob/28869401eff8fcd94aa68d662fdd4d3dbe2f58af/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala
But we could also just use an ultimate `flatMapGroupsWithState` that achieves the same goal as @zhenlineo suggested and still delete this helper method.
@zhenlineo @bogao007 wdyt?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1231405201
##########
connector/connect/common/src/main/java/org/apache/spark/api/java/function/FlatMapGroupsWithStateFunction.java:
##########
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.streaming.GroupState;
+
+/**
+ * ::Experimental::
+ * Base interface for a map function used in
+ * {@code org.apache.spark.sql.KeyValueGroupedDataset.flatMapGroupsWithState(
+ * FlatMapGroupsWithStateFunction, org.apache.spark.sql.streaming.OutputMode,
+ * org.apache.spark.sql.Encoder, org.apache.spark.sql.Encoder)}
+ * @since 2.1.1
Review Comment:
The dependencies are actually in `sql/core` which is not accessible from the common package, so we might have to duplicate these 2 functions. Let me know which util package I should include these functions, otherwise we can leave them as it is and refactor them later when we have a solid framework there.
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on PR #41558:
URL: https://github.com/apache/spark/pull/41558#issuecomment-1600055858
> Seems like tests failed. Mind fixing that? Otherwise, should be good to go I susepct?
Hi @HyukjinKwon, the build succeeded, could you help to merge the code? Thanks!
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227294305
##########
connector/connect/common/src/main/scala/org/apache/spark/sql/connect/common/UdfUtils.scala:
##########
@@ -95,6 +96,26 @@ private[sql] object UdfUtils extends Serializable {
}
}
+ def mapValuesAdaptor[K, V, S, U, IV](
+ f: (K, Iterator[V], GroupState[S]) => U,
+ valueMapFunc: IV => V): (K, Iterator[IV], GroupState[S]) => U = {
+ (k: K, itr: Iterator[IV], s: GroupState[S]) =>
+ {
+ f(k, itr.map(v => valueMapFunc(v)), s)
+ }
+ }
+
+ def mapGroupsWithStateFuncToScalaFunc[K, V, S, U](
+ f: MapGroupsWithStateFunction[K, V, S, U]): (K, Iterator[V], GroupState[S]) => U = {
+ (key, data, groupState) => f.call(key, data.asJava, groupState)
+ }
+
+ def flatMapGroupsWithStateFuncToScalaFunc[K, V, S, U](
Review Comment:
With this function, I think we also need to add `java-specific` APIs to `KeyValueGroupedDataset`?
For example:
https://github.com/apache/spark/blob/28869401eff8fcd94aa68d662fdd4d3dbe2f58af/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala#L118-L136
and this:
https://github.com/apache/spark/blob/28869401eff8fcd94aa68d662fdd4d3dbe2f58af/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala#L575-L603
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] zhenlineo commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "zhenlineo (via GitHub)" <gi...@apache.org>.
zhenlineo commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227339536
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -460,6 +461,159 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
cogroupSorted(other)(thisSortExprs: _*)(otherSortExprs: _*)(
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ mapGroupsWithState(GroupStateTimeout.NoTimeout)(func)
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param timeoutConf
+ * Timeout configuration for groups that do not receive data for a while.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ throw new UnsupportedOperationException
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param timeoutConf
+ * Timeout Conf, see GroupStateTimeout for more details
+ * @param initialState
+ * The user provided state that will be initialized when the first batch of data is processed
+ * in the streaming query. The user defined function will be called on the state data even if
+ * there are no other values in the group. To convert a Dataset ds of type Dataset[(K, S)] to
+ * a KeyValueGroupedDataset[K, S] do {{{ds.groupByKey(x => x._1).mapValues(_._2)}}}
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ throw new UnsupportedOperationException
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See `GroupState` for more
+ * details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param outputMode
+ * The output mode of the function.
+ * @param timeoutConf
+ * Timeout configuration for groups that do not receive data for a while.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ throw new UnsupportedOperationException
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See `GroupState` for more
+ * details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param outputMode
+ * The output mode of the function.
+ * @param timeoutConf
+ * Timeout configuration for groups that do not receive data for a while.
+ * @param initialState
+ * The user provided state that will be initialized when the first batch of data is processed
+ * in the streaming query. The user defined function will be called on the state data even if
+ * there are no other values in the group. To covert a Dataset `ds` of type of type
+ * `Dataset[(K, S)]` to a `KeyValueGroupedDataset[K, S]`, use
+ * {{{ds.groupByKey(x => x._1).mapValues(_._2)}}} See [[Encoder]] for more details on what
+ * types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ throw new UnsupportedOperationException
+ }
+
+ /**
+ * Returns the current java.util.List[proto.Expression].
+ *
+ * @since 3.5.0
+ */
+ def getGroupingExpressions: java.util.List[proto.Expression] = {
Review Comment:
This will go to the public API, do you intend to add this to public API?
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -460,6 +461,159 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
cogroupSorted(other)(thisSortExprs: _*)(otherSortExprs: _*)(
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ mapGroupsWithState(GroupStateTimeout.NoTimeout)(func)
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param timeoutConf
+ * Timeout configuration for groups that do not receive data for a while.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ throw new UnsupportedOperationException
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param timeoutConf
+ * Timeout Conf, see GroupStateTimeout for more details
+ * @param initialState
+ * The user provided state that will be initialized when the first batch of data is processed
+ * in the streaming query. The user defined function will be called on the state data even if
+ * there are no other values in the group. To convert a Dataset ds of type Dataset[(K, S)] to
+ * a KeyValueGroupedDataset[K, S] do {{{ds.groupByKey(x => x._1).mapValues(_._2)}}}
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ throw new UnsupportedOperationException
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See `GroupState` for more
+ * details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param outputMode
+ * The output mode of the function.
+ * @param timeoutConf
+ * Timeout configuration for groups that do not receive data for a while.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def flatMapGroupsWithState[S: Encoder, U: Encoder](
Review Comment:
The only diff between this and the method with the same name is if `initialState` is provided or not. Can you delegate the impl of these two methods to one protected method internally?
```
protected def flatMapGroupsWithState( ..., optionalInitState...)(...) = {
throw new UnsupportedOperationException
}
```
And then impl in the child class? See `aggUntyped` as an example.
##########
connector/connect/common/src/main/protobuf/spark/connect/relations.proto:
##########
@@ -905,6 +906,29 @@ message ApplyInPandasWithState {
string timeout_conf = 7;
}
+message FlatMapGroupsWithState {
Review Comment:
Is it the same as `GroupMap`? Any reason that you cannot reuse the same message?
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -460,6 +461,159 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
cogroupSorted(other)(thisSortExprs: _*)(otherSortExprs: _*)(
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ mapGroupsWithState(GroupStateTimeout.NoTimeout)(func)
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param timeoutConf
+ * Timeout configuration for groups that do not receive data for a while.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
Review Comment:
I feel `mapGroupsWithState` can be expressed using `flatMapGroupsWithState`. Could you delegate the call to `flatMapGroupsWithState` to reduce the methods that need to impl in the child class?
##########
connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/KeyValueGroupedDatasetE2ETestSuite.scala:
##########
@@ -447,4 +449,20 @@ class KeyValueGroupedDatasetE2ETestSuite extends QueryTest with SQLHelper {
checkDataset(keys, "1", "2", "10", "20")
}
+
+ test("flatMapGroupsWithState") {
+ val stateFunc = (key: String, values: Iterator[String], state: GroupState[Int]) => {
+ if (state.exists) throw new IllegalArgumentException("state.exists should be false")
+ Iterator((key, values.size))
+ }
+
+ val session: SparkSession = spark
+ import session.implicits._
+ val values = Seq("a", "a", "b", "c", "c", "c", "c").toDS()
+ .groupByKey(x => x)
+ .flatMapGroupsWithState(Append, GroupStateTimeout.NoTimeout)(stateFunc)
Review Comment:
Feels the conversion func is wrongly constructed. You can check out some conversion code I wrote for `mapGroups -> flatMapGroups`?
```
def mapGroups[U: Encoder](f: (K, Iterator[V]) => U): Dataset[U] = {
flatMapGroups(UdfUtils.mapGroupsFuncToFlatMapAdaptor(f))
}
```
##########
connector/connect/common/src/main/java/org/apache/spark/api/java/function/FlatMapGroupsWithStateFunction.java:
##########
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.streaming.GroupState;
+
+/**
+ * ::Experimental::
+ * Base interface for a map function used in
+ * {@code org.apache.spark.sql.KeyValueGroupedDataset.flatMapGroupsWithState(
+ * FlatMapGroupsWithStateFunction, org.apache.spark.sql.streaming.OutputMode,
+ * org.apache.spark.sql.Encoder, org.apache.spark.sql.Encoder)}
+ * @since 2.1.1
Review Comment:
This funcs need to be moved to the util package. I will leave it to you if you want to directly depends on it or copy the files for now.
In the future we should refactoring with other such classes in the util package.
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1231407237
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -462,6 +462,15 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+ protected def flatMapGroupsWithStateHelper[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S],
Review Comment:
I feel this helper method shouldn't appear there. Is this change because of this comment? https://github.com/apache/spark/pull/41558#discussion_r1227337649
I actually think it's good to use the original implementation, i.e. have a dedicated `flatMapGroupsWithState` and `mapGroupsWIthState` that handles all cases respectively, this also follows the original API design: https://github.com/apache/spark/blob/28869401eff8fcd94aa68d662fdd4d3dbe2f58af/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala
But we could also just use an ultimate `flatMapGroupsWithState` that achieves the same goal as @zhenlineo suggested and still delete this helper method.
Sorry for the back and forth. @zhenlineo @bogao007 wdyt?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227320572
##########
connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/KeyValueGroupedDatasetE2ETestSuite.scala:
##########
@@ -447,4 +449,20 @@ class KeyValueGroupedDatasetE2ETestSuite extends QueryTest with SQLHelper {
checkDataset(keys, "1", "2", "10", "20")
}
+
+ test("flatMapGroupsWithState") {
+ val stateFunc = (key: String, values: Iterator[String], state: GroupState[Int]) => {
+ if (state.exists) throw new IllegalArgumentException("state.exists should be false")
+ Iterator((key, values.size))
+ }
+
+ val session: SparkSession = spark
+ import session.implicits._
+ val values = Seq("a", "a", "b", "c", "c", "c", "c").toDS()
+ .groupByKey(x => x)
+ .flatMapGroupsWithState(Append, GroupStateTimeout.NoTimeout)(stateFunc)
Review Comment:
When I tried to add `mapGroupsWithState` test, it actually failed with below error when I call `.collect()`.
```
io.grpc.StatusRuntimeException: INTERNAL: Job aborted due to stage failure: Task 0 in stage 112.0 failed 1 times, most recent failure: Lost task 0.0 in stage 112.0 (TID 270) (192.168.2.82 executor driver): java.lang.ClassCastException: class scala.Tuple2 cannot be cast to class scala.collection.TraversableOnce (scala.Tuple2 and scala.collection.TraversableOnce are in unnamed module of loader 'app')
```
Do you have any insights on this? Not sure if it's an issue with my implementation or tests.
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon closed pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon closed pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
URL: https://github.com/apache/spark/pull/41558
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on PR #41558:
URL: https://github.com/apache/spark/pull/41558#issuecomment-1593552971
Thanks for the change!
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rangadi commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "rangadi (via GitHub)" <gi...@apache.org>.
rangadi commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227175955
##########
connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/KeyValueGroupedDatasetE2ETestSuite.scala:
##########
@@ -447,4 +449,20 @@ class KeyValueGroupedDatasetE2ETestSuite extends QueryTest with SQLHelper {
checkDataset(keys, "1", "2", "10", "20")
}
+
+ test("flatMapGroupsWithState") {
Review Comment:
Can you try a test with user class?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1230200585
##########
connector/connect/common/src/main/java/org/apache/spark/api/java/function/FlatMapGroupsWithStateFunction.java:
##########
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.streaming.GroupState;
+
+/**
+ * ::Experimental::
+ * Base interface for a map function used in
+ * {@code org.apache.spark.sql.KeyValueGroupedDataset.flatMapGroupsWithState(
+ * FlatMapGroupsWithStateFunction, org.apache.spark.sql.streaming.OutputMode,
+ * org.apache.spark.sql.Encoder, org.apache.spark.sql.Encoder)}
+ * @since 2.1.1
Review Comment:
The funcs need to be under common package since it's also being used by `UdfUtils`. I don't see a util package in common, do you want me to create one? I think previously you mentioned that there was an effort to build that util package and framework, is that already completed?
##########
connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/KeyValueGroupedDatasetE2ETestSuite.scala:
##########
@@ -447,4 +449,20 @@ class KeyValueGroupedDatasetE2ETestSuite extends QueryTest with SQLHelper {
checkDataset(keys, "1", "2", "10", "20")
}
+
+ test("flatMapGroupsWithState") {
+ val stateFunc = (key: String, values: Iterator[String], state: GroupState[Int]) => {
+ if (state.exists) throw new IllegalArgumentException("state.exists should be false")
+ Iterator((key, values.size))
+ }
+
+ val session: SparkSession = spark
+ import session.implicits._
+ val values = Seq("a", "a", "b", "c", "c", "c", "c").toDS()
+ .groupByKey(x => x)
+ .flatMapGroupsWithState(Append, GroupStateTimeout.NoTimeout)(stateFunc)
Review Comment:
Thanks @zhenlineo, it worked after updating the conversion func!
##########
connector/connect/common/src/main/scala/org/apache/spark/sql/connect/common/UdfUtils.scala:
##########
@@ -95,6 +96,26 @@ private[sql] object UdfUtils extends Serializable {
}
}
+ def mapValuesAdaptor[K, V, S, U, IV](
+ f: (K, Iterator[V], GroupState[S]) => U,
+ valueMapFunc: IV => V): (K, Iterator[IV], GroupState[S]) => U = {
+ (k: K, itr: Iterator[IV], s: GroupState[S]) =>
+ {
+ f(k, itr.map(v => valueMapFunc(v)), s)
+ }
+ }
+
+ def mapGroupsWithStateFuncToScalaFunc[K, V, S, U](
+ f: MapGroupsWithStateFunction[K, V, S, U]): (K, Iterator[V], GroupState[S]) => U = {
+ (key, data, groupState) => f.call(key, data.asJava, groupState)
+ }
+
+ def flatMapGroupsWithStateFuncToScalaFunc[K, V, S, U](
Review Comment:
Added jave specific ones
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on PR #41558:
URL: https://github.com/apache/spark/pull/41558#issuecomment-1593494039
@WweiL Please help to take another look, thanks!
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227308844
##########
connector/connect/common/src/main/scala/org/apache/spark/sql/connect/common/UdfUtils.scala:
##########
@@ -95,6 +96,26 @@ private[sql] object UdfUtils extends Serializable {
}
}
+ def mapValuesAdaptor[K, V, S, U, IV](
+ f: (K, Iterator[V], GroupState[S]) => U,
+ valueMapFunc: IV => V): (K, Iterator[IV], GroupState[S]) => U = {
+ (k: K, itr: Iterator[IV], s: GroupState[S]) =>
+ {
+ f(k, itr.map(v => valueMapFunc(v)), s)
+ }
+ }
+
+ def mapGroupsWithStateFuncToScalaFunc[K, V, S, U](
+ f: MapGroupsWithStateFunction[K, V, S, U]): (K, Iterator[V], GroupState[S]) => U = {
+ (key, data, groupState) => f.call(key, data.asJava, groupState)
+ }
+
+ def flatMapGroupsWithStateFuncToScalaFunc[K, V, S, U](
Review Comment:
Will include those, thanks!
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -572,6 +726,101 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV](
agg(aggregator)
}
+ override def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
Review Comment:
This actually failed the scalastyle check, will run the scalastyle check locally, that should fix all these kind of formatting issues
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on PR #41558:
URL: https://github.com/apache/spark/pull/41558#issuecomment-1588217946
CC @zhenlineo Can you also take a look at this PR when you get a chance? Thank you so much!
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227312489
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -572,6 +726,101 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV](
agg(aggregator)
}
+ override def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def getGroupingExpressions: java.util.List[proto.Expression] = {
Review Comment:
If my understanding is correct, this function servers as a public getter of `groupingExprs` so that `initialState.groupingExprs` can be called.
Can we follow the implementation here:
https://github.com/apache/spark/blob/db3f85b24498d4904134ed8f14b570fde6550060/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala#L685-L703
To get rid of this function completely?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] zhenlineo commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "zhenlineo (via GitHub)" <gi...@apache.org>.
zhenlineo commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1231397503
##########
connector/connect/common/src/main/java/org/apache/spark/api/java/function/FlatMapGroupsWithStateFunction.java:
##########
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.java.function;
+
+import java.io.Serializable;
+import java.util.Iterator;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.streaming.GroupState;
+
+/**
+ * ::Experimental::
+ * Base interface for a map function used in
+ * {@code org.apache.spark.sql.KeyValueGroupedDataset.flatMapGroupsWithState(
+ * FlatMapGroupsWithStateFunction, org.apache.spark.sql.streaming.OutputMode,
+ * org.apache.spark.sql.Encoder, org.apache.spark.sql.Encoder)}
+ * @since 2.1.1
Review Comment:
common currently depends on core and catalyst, so you should have your dependencies there too? In the future, common will depends on some `util` package, which should contains these two classes of yours.
If yours is in core, catalyst, you can leave the refactoring to us later.
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227307156
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -460,6 +461,159 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
cogroupSorted(other)(thisSortExprs: _*)(otherSortExprs: _*)(
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ mapGroupsWithState(GroupStateTimeout.NoTimeout)(func)
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param timeoutConf
+ * Timeout configuration for groups that do not receive data for a while.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
Review Comment:
nvm see your comments here https://github.com/apache/spark/pull/40959#discussion_r1206059114
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227319913
##########
connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/KeyValueGroupedDatasetE2ETestSuite.scala:
##########
Review Comment:
Also can you add a test when there are initial states?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] rangadi commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "rangadi (via GitHub)" <gi...@apache.org>.
rangadi commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227271011
##########
connector/connect/common/src/main/protobuf/spark/connect/relations.proto:
##########
@@ -905,6 +906,29 @@ message ApplyInPandasWithState {
string timeout_conf = 7;
}
+message FlatMapGroupsWithState {
Review Comment:
Add documentation
##########
connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/KeyValueGroupedDatasetE2ETestSuite.scala:
##########
@@ -447,4 +449,20 @@ class KeyValueGroupedDatasetE2ETestSuite extends QueryTest with SQLHelper {
checkDataset(keys, "1", "2", "10", "20")
}
+
+ test("flatMapGroupsWithState") {
+ val stateFunc = (key: String, values: Iterator[String], state: GroupState[Int]) => {
+ if (state.exists) throw new IllegalArgumentException("state.exists should be false")
+ Iterator((key, values.size))
+ }
+
+ val session: SparkSession = spark
+ import session.implicits._
+ val values = Seq("a", "a", "b", "c", "c", "c", "c").toDS()
+ .groupByKey(x => x)
+ .flatMapGroupsWithState(Append, GroupStateTimeout.NoTimeout)(stateFunc)
Review Comment:
Could you add a test for mapGroupsWithState() as well? It could be similarly simple.
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227301618
##########
connector/connect/common/src/main/scala/org/apache/spark/sql/streaming/GroupState.scala:
##########
@@ -0,0 +1,336 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.streaming
+
+import org.apache.spark.annotation.{Evolving, Experimental}
+import org.apache.spark.sql.catalyst.plans.logical.LogicalGroupState
+
+/**
+ * :: Experimental ::
+ *
+ * Wrapper class for interacting with per-group state data in `mapGroupsWithState` and
+ * `flatMapGroupsWithState` operations on `KeyValueGroupedDataset`.
+ *
+ * Detail description on `[map/flatMap]GroupsWithState` operation
+ * -------------------------------------------------------------- Both, `mapGroupsWithState` and
+ * `flatMapGroupsWithState` in `KeyValueGroupedDataset` will invoke the user-given function on
+ * each group (defined by the grouping function in `Dataset.groupByKey()`) while maintaining a
+ * user-defined per-group state between invocations. For a static batch Dataset, the function will
+ * be invoked once per group. For a streaming Dataset, the function will be invoked for each group
+ * repeatedly in every trigger. That is, in every batch of the `StreamingQuery`, the function will
+ * be invoked once for each group that has data in the trigger. Furthermore, if timeout is set,
+ * then the function will be invoked on timed-out groups (more detail below).
+ *
+ * The function is invoked with the following parameters.
+ * - The key of the group.
+ * - An iterator containing all the values for this group.
+ * - A user-defined state object set by previous invocations of the given function.
+ *
+ * In case of a batch Dataset, there is only one invocation and the state object will be empty as
+ * there is no prior state. Essentially, for batch Datasets, `[map/flatMap]GroupsWithState` is
+ * equivalent to `[map/flatMap]Groups` and any updates to the state and/or timeouts have no
+ * effect.
+ *
+ * The major difference between `mapGroupsWithState` and `flatMapGroupsWithState` is that the
+ * former allows the function to return one and only one record, whereas the latter allows the
+ * function to return any number of records (including no records). Furthermore, the
+ * `flatMapGroupsWithState` is associated with an operation output mode, which can be either
+ * `Append` or `Update`. Semantically, this defines whether the output records of one trigger is
+ * effectively replacing the previously output records (from previous triggers) or is appending to
+ * the list of previously output records. Essentially, this defines how the Result Table (refer to
+ * the semantics in the programming guide) is updated, and allows us to reason about the semantics
+ * of later operations.
+ *
+ * Important points to note about the function (both mapGroupsWithState and
+ * flatMapGroupsWithState).
+ * - In a trigger, the function will be called only the groups present in the batch. So do not
+ * assume that the function will be called in every trigger for every group that has state.
+ * - There is no guaranteed ordering of values in the iterator in the function, neither with
+ * batch, nor with streaming Datasets.
+ * - All the data will be shuffled before applying the function.
+ * - If timeout is set, then the function will also be called with no values. See more details
+ * on `GroupStateTimeout` below.
+ *
+ * Important points to note about using `GroupState`.
+ * - The value of the state cannot be null. So updating state with null will throw
+ * `IllegalArgumentException`.
+ * - Operations on `GroupState` are not thread-safe. This is to avoid memory barriers.
+ * - If `remove()` is called, then `exists()` will return `false`, `get()` will throw
+ * `NoSuchElementException` and `getOption()` will return `None`
+ * - After that, if `update(newState)` is called, then `exists()` will again return `true`,
+ * `get()` and `getOption()`will return the updated value.
+ *
+ * Important points to note about using `GroupStateTimeout`.
+ * - The timeout type is a global param across all the groups (set as `timeout` param in
+ * `[map|flatMap]GroupsWithState`, but the exact timeout duration/timestamp is configurable
+ * per group by calling `setTimeout...()` in `GroupState`.
+ * - Timeouts can be either based on processing time (i.e.
+ * `GroupStateTimeout.ProcessingTimeTimeout`) or event time (i.e.
+ * `GroupStateTimeout.EventTimeTimeout`).
+ * - With `ProcessingTimeTimeout`, the timeout duration can be set by calling
+ * `GroupState.setTimeoutDuration`. The timeout will occur when the clock has advanced by the
+ * set duration. Guarantees provided by this timeout with a duration of D ms are as follows:
+ * - Timeout will never occur before the clock time has advanced by D ms
+ * - Timeout will occur eventually when there is a trigger in the query (i.e. after D ms). So
+ * there is no strict upper bound on when the timeout would occur. For example, the trigger
+ * interval of the query will affect when the timeout actually occurs. If there is no data
+ * in the stream (for any group) for a while, then there will not be any trigger and timeout
+ * function call will not occur until there is data.
+ * - Since the processing time timeout is based on the clock time, it is affected by the
+ * variations in the system clock (i.e. time zone changes, clock skew, etc.).
+ * - With `EventTimeTimeout`, the user also has to specify the event time watermark in the query
+ * using `Dataset.withWatermark()`. With this setting, data that is older than the watermark
+ * is filtered out. The timeout can be set for a group by setting a timeout timestamp
+ * using`GroupState.setTimeoutTimestamp()`, and the timeout would occur when the watermark
+ * advances beyond the set timestamp. You can control the timeout delay by two parameters -
+ * (i) watermark delay and an additional duration beyond the timestamp in the event (which is
+ * guaranteed to be newer than watermark due to the filtering). Guarantees provided by this
+ * timeout are as follows:
+ * - Timeout will never occur before the watermark has exceeded the set timeout.
+ * - Similar to processing time timeouts, there is no strict upper bound on the delay when the
+ * timeout actually occurs. The watermark can advance only when there is data in the stream
+ * and the event time of the data has actually advanced.
+ * - When the timeout occurs for a group, the function is called for that group with no values,
+ * and `GroupState.hasTimedOut()` set to true.
+ * - The timeout is reset every time the function is called on a group, that is, when the group
+ * has new data, or the group has timed out. So the user has to set the timeout duration every
+ * time the function is called, otherwise, there will not be any timeout set.
+ *
+ * `[map/flatMap]GroupsWithState` can take a user defined initial state as an additional argument.
+ * This state will be applied when the first batch of the streaming query is processed. If there
+ * are no matching rows in the data for the keys present in the initial state, the state is still
+ * applied and the function will be invoked with the values being an empty iterator.
+ *
+ * Scala example of using GroupState in `mapGroupsWithState`:
+ * {{{
+ * // A mapping function that maintains an integer state for string keys and returns a string.
+ * // Additionally, it sets a timeout to remove the state if it has not received data for an hour.
+ * def mappingFunction(key: String, value: Iterator[Int], state: GroupState[Int]): String = {
+ *
+ * if (state.hasTimedOut) { // If called when timing out, remove the state
+ * state.remove()
+ *
+ * } else if (state.exists) { // If state exists, use it for processing
+ * val existingState = state.get // Get the existing state
+ * val shouldRemove = ... // Decide whether to remove the state
+ * if (shouldRemove) {
+ * state.remove() // Remove the state
+ *
+ * } else {
+ * val newState = ...
+ * state.update(newState) // Set the new state
+ * state.setTimeoutDuration("1 hour") // Set the timeout
+ * }
+ *
+ * } else {
+ * val initialState = ...
+ * state.update(initialState) // Set the initial state
+ * state.setTimeoutDuration("1 hour") // Set the timeout
+ * }
+ * ...
+ * // return something
+ * }
+ *
+ * dataset
+ * .groupByKey(...)
+ * .mapGroupsWithState(GroupStateTimeout.ProcessingTimeTimeout)(mappingFunction)
+ * }}}
+ *
+ * Java example of using `GroupState`:
+ * {{{
+ * // A mapping function that maintains an integer state for string keys and returns a string.
+ * // Additionally, it sets a timeout to remove the state if it has not received data for an hour.
+ * MapGroupsWithStateFunction<String, Integer, Integer, String> mappingFunction =
+ * new MapGroupsWithStateFunction<String, Integer, Integer, String>() {
+ *
+ * @Override
+ * public String call(String key, Iterator<Integer> value, GroupState<Integer> state) {
+ * if (state.hasTimedOut()) { // If called when timing out, remove the state
+ * state.remove();
+ *
+ * } else if (state.exists()) { // If state exists, use it for processing
+ * int existingState = state.get(); // Get the existing state
+ * boolean shouldRemove = ...; // Decide whether to remove the state
+ * if (shouldRemove) {
+ * state.remove(); // Remove the state
+ *
+ * } else {
+ * int newState = ...;
+ * state.update(newState); // Set the new state
+ * state.setTimeoutDuration("1 hour"); // Set the timeout
+ * }
+ *
+ * } else {
+ * int initialState = ...; // Set the initial state
+ * state.update(initialState);
+ * state.setTimeoutDuration("1 hour"); // Set the timeout
+ * }
+ * ...
+ * // return something
+ * }
+ * };
+ *
+ * dataset
+ * .groupByKey(...)
+ * .mapGroupsWithState(
+ * mappingFunction, Encoders.INT, Encoders.STRING, GroupStateTimeout.ProcessingTimeTimeout);
+ * }}}
+ *
+ * @tparam S
+ * User-defined type of the state to be stored for each group. Must be encodable into Spark SQL
+ * types (see `Encoder` for more details).
+ * @since 2.2.0
Review Comment:
nit version
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227322575
##########
connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/KeyValueGroupedDatasetE2ETestSuite.scala:
##########
@@ -447,4 +449,20 @@ class KeyValueGroupedDatasetE2ETestSuite extends QueryTest with SQLHelper {
checkDataset(keys, "1", "2", "10", "20")
}
+
+ test("flatMapGroupsWithState") {
+ val stateFunc = (key: String, values: Iterator[String], state: GroupState[Int]) => {
+ if (state.exists) throw new IllegalArgumentException("state.exists should be false")
+ Iterator((key, values.size))
+ }
+
+ val session: SparkSession = spark
+ import session.implicits._
+ val values = Seq("a", "a", "b", "c", "c", "c", "c").toDS()
+ .groupByKey(x => x)
+ .flatMapGroupsWithState(Append, GroupStateTimeout.NoTimeout)(stateFunc)
Review Comment:
cc @zhenlineo
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227312489
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -572,6 +726,101 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV](
agg(aggregator)
}
+ override def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def getGroupingExpressions: java.util.List[proto.Expression] = {
Review Comment:
If my understand is correct, this function servers as a public getter of `groupingExprs` so that `initialState.groupingExprs` can be called.
Can we follow the implementation here:
https://github.com/apache/spark/blob/db3f85b24498d4904134ed8f14b570fde6550060/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala#L685-L703
To get rid of this function completely?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227284034
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -460,6 +461,159 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
cogroupSorted(other)(thisSortExprs: _*)(otherSortExprs: _*)(
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ mapGroupsWithState(GroupStateTimeout.NoTimeout)(func)
+ }
+
+ /**
+ * (Scala-specific) Applies the given function to each group of data, while maintaining a
+ * user-defined per-group state. The result Dataset will represent the objects returned by the
+ * function. For a static batch Dataset, the function will be invoked once per group. For a
+ * streaming Dataset, the function will be invoked for each group repeatedly in every trigger,
+ * and updates to each group's state will be saved across invocations. See
+ * [[org.apache.spark.sql.streaming.GroupState]] for more details.
+ *
+ * @tparam S
+ * The type of the user-defined state. Must be encodable to Spark SQL types.
+ * @tparam U
+ * The type of the output objects. Must be encodable to Spark SQL types.
+ * @param func
+ * Function to be called on every group.
+ * @param timeoutConf
+ * Timeout configuration for groups that do not receive data for a while.
+ *
+ * See [[Encoder]] for more details on what types are encodable to Spark SQL.
+ * @since 3.5.0
+ */
+ def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
Review Comment:
nit: indentation
Also may I know why is this and the following functions not supported?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227312489
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -572,6 +726,101 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV](
agg(aggregator)
}
+ override def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def getGroupingExpressions: java.util.List[proto.Expression] = {
Review Comment:
If my understand is correct, this function exists for a public reader of `groupingExprs` so that `initialState.groupingExprs` can be called.
Can we follow the implementation here:
https://github.com/apache/spark/blob/db3f85b24498d4904134ed8f14b570fde6550060/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala#L685-L703
To get rid of this function completely?
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -572,6 +726,101 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV](
agg(aggregator)
}
+ override def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def getGroupingExpressions: java.util.List[proto.Expression] = {
Review Comment:
If my understand is correct, this function servers as a public reader of `groupingExprs` so that `initialState.groupingExprs` can be called.
Can we follow the implementation here:
https://github.com/apache/spark/blob/db3f85b24498d4904134ed8f14b570fde6550060/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala#L685-L703
To get rid of this function completely?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1231394671
##########
connector/connect/common/src/main/scala/org/apache/spark/sql/connect/common/UdfUtils.scala:
##########
@@ -95,6 +96,31 @@ private[sql] object UdfUtils extends Serializable {
}
}
+ def mapValuesAdaptor[K, V, S, U, IV](
+ f: (K, Iterator[V], GroupState[S]) => Iterator[U],
+ valueMapFunc: IV => V): (K, Iterator[IV], GroupState[S]) => Iterator[U] = {
+ print("Used function: mapValuesAdaptor")
Review Comment:
Good catch! Forgot to remove it...
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] zhenlineo commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "zhenlineo (via GitHub)" <gi...@apache.org>.
zhenlineo commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1231388771
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -572,6 +916,55 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV](
agg(aggregator)
}
+ override protected def flatMapGroupsWithStateHelper[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S],
+ isMapGroupWithState: Boolean)(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != null && outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
Review Comment:
nit: How about use `Option` for the nullable fields?
##########
connector/connect/common/src/main/scala/org/apache/spark/sql/connect/common/UdfUtils.scala:
##########
@@ -95,6 +96,31 @@ private[sql] object UdfUtils extends Serializable {
}
}
+ def mapValuesAdaptor[K, V, S, U, IV](
+ f: (K, Iterator[V], GroupState[S]) => Iterator[U],
+ valueMapFunc: IV => V): (K, Iterator[IV], GroupState[S]) => Iterator[U] = {
+ print("Used function: mapValuesAdaptor")
Review Comment:
Debugging code?
##########
connector/connect/common/src/main/protobuf/spark/connect/relations.proto:
##########
@@ -857,6 +857,21 @@ message GroupMap {
// (Optional) Expressions for sorting. Only used by Scala Sorted Group Map API.
repeated Expression sorting_expressions = 4;
+
+ // (Optional) Input relation for initial State.
+ Relation initial_input = 5;
+
+ // (Optional) Expressions for grouping keys of the initial state input relation.
+ repeated Expression initial_grouping_expressions = 6;
+
+ // (Optional) True if MapGroupsWithState, false if FlatMapGroupsWithState.
+ optional bool is_map_groups_with_state = 7;
+
+ // (Optional) The output mode of the function.
+ optional string output_mode = 8;
+
+ // (Optional) Timeout configuration for groups that do not receive data for a while.
+ optional string timeout_conf = 9;
Review Comment:
cc @grundprinzip @xinrong-meng for the proto change.
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1231423248
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -462,6 +462,15 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+ protected def flatMapGroupsWithStateHelper[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S],
Review Comment:
> But we could also just use an ultimate `flatMapGroupsWithState` that achieves the same goal as @zhenlineo suggested and still delete this helper method.
It would be hard to achieve this since the original `flatMapGroupsWithState` methods doesn't contain a boolean to decide if it's `flatMapGroupsWithState` or `mapGroupsWIthState`. If we add that boolean, then it's the same as the helper function approach. Otherwise we have to maintain 2 different implementations
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -462,6 +462,15 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+ protected def flatMapGroupsWithStateHelper[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S],
Review Comment:
> But we could also just use an ultimate `flatMapGroupsWithState` that achieves the same goal as @zhenlineo suggested and still delete this helper method.
It would be hard to achieve this since the original `flatMapGroupsWithState` methods doesn't contain a boolean to decide if it's `flatMapGroupsWithState` or `mapGroupsWIthState`. If we add that boolean, then it's the same as the helper function approach. Otherwise we have to maintain 2 different implementations
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227312489
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -572,6 +726,101 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV](
agg(aggregator)
}
+ override def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def getGroupingExpressions: java.util.List[proto.Expression] = {
Review Comment:
Can we follow the implementation here:
https://github.com/apache/spark/blob/db3f85b24498d4904134ed8f14b570fde6550060/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala#L685-L703
To get rid of this function completely?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227396198
##########
connector/connect/common/src/main/protobuf/spark/connect/relations.proto:
##########
@@ -905,6 +906,29 @@ message ApplyInPandasWithState {
string timeout_conf = 7;
}
+message FlatMapGroupsWithState {
Review Comment:
If I want to reuse the existing `GroupMap`, do I just add new fields as optional? Also, what's the recommend way to add optional fields? I saw that in `GroupMap`, `sorting_expressions` is commented as optional, but it doesn't have an `optional` keyword in front.
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] zhenlineo commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "zhenlineo (via GitHub)" <gi...@apache.org>.
zhenlineo commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1231414703
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -462,6 +462,15 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+ protected def flatMapGroupsWithStateHelper[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S],
Review Comment:
> But we could also just use an ultimate flatMapGroupsWithState that achieves the same goal as @zhenlineo suggested and still delete this helper method.
Fine with me. I mostly want to keep the `Impl` class small as the logics there is already so complicated.
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] zhenlineo commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "zhenlineo (via GitHub)" <gi...@apache.org>.
zhenlineo commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1228257497
##########
connector/connect/common/src/main/protobuf/spark/connect/relations.proto:
##########
@@ -905,6 +906,29 @@ message ApplyInPandasWithState {
string timeout_conf = 7;
}
+message FlatMapGroupsWithState {
Review Comment:
In short, all fields in grpc3 are optional. When explicitly adding `optional` keywords, it add a `hasXXX` boolean method. If you need hasXXX, you need to prefix the `optional` keyword.
For the complete story, you can refer to the [doc](https://protobuf.dev/programming-guides/proto3/#field-labels).
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] WweiL commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "WweiL (via GitHub)" <gi...@apache.org>.
WweiL commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227312489
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -572,6 +726,101 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV](
agg(aggregator)
}
+ override def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def getGroupingExpressions: java.util.List[proto.Expression] = {
Review Comment:
If my understand is correct, this function servers as a public getter of `groupingExprs` so that `initialState.groupingExprs` can be called above.
Can we follow the implementation here:
https://github.com/apache/spark/blob/db3f85b24498d4904134ed8f14b570fde6550060/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala#L685-L703
To get rid of this function completely?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227319087
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -572,6 +726,101 @@ private class KeyValueGroupedDatasetImpl[K, V, IK, IV](
agg(aggregator)
}
+ override def mapGroupsWithState[S: Encoder, U: Encoder](
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def mapGroupsWithState[S: Encoder, U: Encoder](timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => U): Dataset[U] = {
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(true)
+ .setOutputMode(OutputMode.Update.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S])(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ .addAllInitialGroupingExpressions(initialState.getGroupingExpressions)
+ }
+ }
+
+ override def flatMapGroupsWithState[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout)(
+ func: (K, Iterator[V], GroupState[S]) => Iterator[U]): Dataset[U] = {
+ if (outputMode != OutputMode.Append && outputMode != OutputMode.Update) {
+ throw new IllegalArgumentException("The output mode of function should be append or update")
+ }
+ val nf = if (valueMapFunc == UdfUtils.identical()) {
+ func
+ } else {
+ UdfUtils.mapValuesAdaptor(func, valueMapFunc)
+ }
+ val outputEncoder = encoderFor[U]
+ sparkSession.newDataset[U](outputEncoder) { builder =>
+ builder.getFlatMapGroupsWithStateBuilder
+ .setInput(plan.getRoot)
+ .addAllGroupingExpressions(groupingExprs)
+ .setFunc(getUdf(nf, outputEncoder)(ivEncoder))
+ .setIsMapGroupsWithState(false)
+ .setOutputMode(outputMode.toString)
+ .setTimeoutConf(timeoutConf.toString)
+ }
+ }
+
+ override def getGroupingExpressions: java.util.List[proto.Expression] = {
Review Comment:
Thanks for the pointer, yeah this is a better way to handle it!
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1227320740
##########
connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/KeyValueGroupedDatasetE2ETestSuite.scala:
##########
Review Comment:
Sure
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #41558:
URL: https://github.com/apache/spark/pull/41558#issuecomment-1596387579
Seems like tests failed. Mind fixing that? Otherwise, should be good to go I susepct?
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] bogao007 commented on a diff in pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #41558:
URL: https://github.com/apache/spark/pull/41558#discussion_r1231419935
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala:
##########
@@ -462,6 +462,15 @@ abstract class KeyValueGroupedDataset[K, V] private[sql] () extends Serializable
UdfUtils.coGroupFunctionToScalaFunc(f))(encoder)
}
+ protected def flatMapGroupsWithStateHelper[S: Encoder, U: Encoder](
+ outputMode: OutputMode,
+ timeoutConf: GroupStateTimeout,
+ initialState: KeyValueGroupedDataset[K, S],
Review Comment:
This is the comment for the change here https://github.com/apache/spark/pull/41558#discussion_r1227341997, we actually already have something similar that's introduced by @zhenlineo [here](https://github.com/apache/spark/blob/49dc2dbcbb1d763d8ee0f2034e82347eebe50d0f/sql/core/src/main/scala/org/apache/spark/sql/KeyValueGroupedDataset.scala#L672-L681).
I would actually prefer the current approach with a helper function since it reduces the code duplication a lot. There are several differences between the original implementation and Spark Connect, a major one is that we separate the implementation out of the public methods, in that case I think a single helper method should be enough.
Another thing is that with the original API design, it doesn't contain much logics in the actual implementation, it just constructs a new `Dataset` object which doesn't make much sense to abstract out any common logics. But for the Spark Connect implementation, it contains some logics there to determine which value we should pass and might result in duplicated code if we have 2 different methods for `flatMapGroupsWithState` and `mapGroupsWIthState`. @WweiL let me know what you think
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on pull request #41558: [SPARK-43511][CONNECT][SS]Implemented MapGroupsWithState and FlatMapGroupsWithState APIs for Spark Connect
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #41558:
URL: https://github.com/apache/spark/pull/41558#issuecomment-1600263691
Merged to master.
--
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.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org