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 2024/08/29 22:30:04 UTC
[PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
bogao007 opened a new pull request, #47933:
URL: https://github.com/apache/spark/pull/47933
<!--
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
'common/utils/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
Support ListState for TransformWithStateInPandas
### Why are the changes needed?
Adding new functionality for TransformWithStateInPandas
### 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 new unit tests.
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has been used in the process of authoring this patch, please include the
phrase: 'Generated-by: ' followed by the name of the tool and its version.
If no, write 'No'.
Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
No
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on PR #47933:
URL: https://github.com/apache/spark/pull/47933#issuecomment-2334758470
Test failure seems related. Could you take a look?
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746098361
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -124,6 +152,8 @@ class TransformWithStateInPandasStateServer(
// The key row is serialized as a byte array, we need to convert it back to a Row
val keyRow = PythonSQLUtils.toJVMRow(keyBytes, groupingKeySchema, keyRowDeserializer)
ImplicitGroupingKeyTracker.setImplicitKey(keyRow)
+ // Reset the list state iterators for a new grouping key.
+ listStateIterators = new mutable.HashMap[String, Iterator[Row]]()
sendResponse(0)
case ImplicitGroupingKeyRequest.MethodCase.REMOVEIMPLICITKEY =>
ImplicitGroupingKeyTracker.removeImplicitKey()
Review Comment:
We should probably also do `listStateIterators = new mutable.HashMap[String, Iterator[Row]]()` here.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -62,6 +74,22 @@ class TransformWithStateInPandasStateServer(
new mutable.HashMap[String, (ValueState[Row], StructType,
ExpressionEncoder.Deserializer[Row])]()
}
+ // A map to store the list state name -> (list state, schema, list state row deserializer,
+ // list state row serializer) mapping.
+ private val listStates = if (listStatesMapForTest != null) {
+ listStatesMapForTest
+ } else {
+ new mutable.HashMap[String, (ListState[Row], StructType,
Review Comment:
Shall we placed the value Tuple into a case class for better readability? Sometimes I need to scroll up to find which index refers to which field here.
##########
python/pyspark/sql/pandas/serializers.py:
##########
@@ -1190,3 +1190,27 @@ def dump_stream(self, iterator, stream):
"""
result = [(b, t) for x in iterator for y, t in x for b in y]
super().dump_stream(result, stream)
+
+
+class TransformWithStateInPandasStateSerializer:
Review Comment:
Maybe I missed something, I don't see this class is called anywhere in the PR?
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1770700243
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator[Row]:
+ """
+ Get list state with an iterator.
+ """
+ return ListStateIterator(self._list_state_client, self._state_name)
+
+ def put(self, new_state: List[Any]) -> None:
Review Comment:
Yeah good to be consistent and Tuple is stricter (more helpful).
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1775665288
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -50,19 +50,11 @@ def exists(self) -> bool:
"""
return self._value_state_client.exists(self._state_name)
- def get(self) -> Optional[Row]:
+ def get(self) -> Optional[Tuple]:
"""
Get the state value if it exists. Returns None if the state variable does not have a value.
"""
- value = self._value_state_client.get(self._state_name)
- if value is None:
- return None
- schema = self.schema
- if isinstance(schema, str):
- schema = cast(StructType, _parse_datatype_string(schema))
- # Create the Row using the values and schema fields
- row = _create_row(schema.fieldNames(), value)
- return row
+ return self._value_state_client.get(self._state_name)
def update(self, new_value: Any) -> None:
Review Comment:
Sure, will address this in the next PR, 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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746332960
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -62,6 +74,22 @@ class TransformWithStateInPandasStateServer(
new mutable.HashMap[String, (ValueState[Row], StructType,
ExpressionEncoder.Deserializer[Row])]()
}
+ // A map to store the list state name -> (list state, schema, list state row deserializer,
+ // list state row serializer) mapping.
+ private val listStates = if (listStatesMapForTest != null) {
+ listStatesMapForTest
+ } else {
+ new mutable.HashMap[String, (ListState[Row], StructType,
+ ExpressionEncoder.Deserializer[Row], ExpressionEncoder.Serializer[Row])]()
+ }
+ // A map to store the list state name -> iterator mapping. This is to keep track of the
+ // current iterator position for each list state in a grouping key in case user tries to fetch
+ // another list state before the current iterator is exhausted.
+ private var listStateIterators = if (listStateIteratorMapForTest != null) {
+ listStateIteratorMapForTest
+ } else {
+ new mutable.HashMap[String, Iterator[Row]]()
Review Comment:
Yeah I think we don't need to worry about thread-safety here since the server thread is per batch and per partition. There would only be one single thread (the server thread itself) accessing the map during the lifecycle. The map would not be accessible from other threads. Let me know if this answers your question.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1750761105
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator["PandasDataFrameLike"]:
Review Comment:
@HeartSaVioR OK, if we'd like to use `Iterator[Row]` in the interfaces, I think it's better to use the same for both `put()` and `get()` instead of `get()` only. Let me see how I could adjust that.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746097930
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -124,6 +152,8 @@ class TransformWithStateInPandasStateServer(
// The key row is serialized as a byte array, we need to convert it back to a Row
val keyRow = PythonSQLUtils.toJVMRow(keyBytes, groupingKeySchema, keyRowDeserializer)
ImplicitGroupingKeyTracker.setImplicitKey(keyRow)
+ // Reset the list state iterators for a new grouping key.
+ listStateIterators = new mutable.HashMap[String, Iterator[Row]]()
sendResponse(0)
case ImplicitGroupingKeyRequest.MethodCase.REMOVEIMPLICITKEY =>
ImplicitGroupingKeyTracker.removeImplicitKey()
Review Comment:
We should probably also do listStateIterators = new mutable.HashMap[String, Iterator[Row]]() here.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1773807929
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -99,25 +99,25 @@ def exists(self) -> bool:
"""
return self._list_state_client.exists(self._state_name)
- def get(self) -> Iterator[Row]:
+ def get(self) -> Iterator[Tuple]:
Review Comment:
Yep, we store as `Row` in state store. The pickling thing is that I followed what ApplyInPandasWithStateWriter does for groupState. It serializes `Row` state to bytes in JVM: https://github.com/apache/spark/blob/55d0233d19cc52bee91a9619057d9b6f33165a0a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ApplyInPandasWithStateWriter.scala#L182 and then deserialize using pickling in python: https://github.com/apache/spark/blob/55d0233d19cc52bee91a9619057d9b6f33165a0a/python/pyspark/sql/pandas/serializers.py#L837 Are you saying we shouldn't use `pickleSer.loads` in valueState `get()`? If so, could you share your concerns here?
An alternative I can think about is to let valueState also use arrow to transmit the state row from JVM to python side and deserialize just like list state. Let me know your thoughts on this, 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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1768129788
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator[Row]:
+ """
+ Get list state with an iterator.
+ """
+ return ListStateIterator(self._list_state_client, self._state_name)
+
+ def put(self, new_state: List[Any]) -> None:
Review Comment:
I understand the type is just a hint, but I wonder the reason we use `Any` and `Row` here. Even further we use `Tuple` in other side. What is the actual type of this? Do we allow multiple types here?
##########
python/pyspark/sql/streaming/stateful_processor_api_client.py:
##########
@@ -168,3 +207,43 @@ def _serialize_to_bytes(self, schema: StructType, data: Tuple) -> bytes:
def _deserialize_from_bytes(self, value: bytes) -> Any:
return self.pickleSer.loads(value)
+
+ def _send_arrow_state(self, schema: StructType, state: List[Tuple]) -> None:
Review Comment:
Is there any utility class/function we can reuse in PySpark? I feel like this is the functionality we should have done before to support Pandas API on Spark. Among `_send_arrow_state` and `_convert`, `_to_numpy_type`.
cc. @HyukjinKwon
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -125,9 +149,13 @@ class TransformWithStateInPandasStateServer(
// The key row is serialized as a byte array, we need to convert it back to a Row
val keyRow = PythonSQLUtils.toJVMRow(keyBytes, groupingKeySchema, keyRowDeserializer)
ImplicitGroupingKeyTracker.setImplicitKey(keyRow)
+ // Reset the list state iterators for a new grouping key.
+ listStateIterators = new mutable.HashMap[String, Iterator[Row]]()
Review Comment:
How do we handle the case of given param `listStateIteratorMapForTest`? Are we assuming the test should only deal with a single grouping key? If then let's explicitly mention in somewhere, e.g. class doc.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServerSuite.scala:
##########
@@ -127,24 +151,79 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
}
test("value state clear") {
- val message = ValueStateCall.newBuilder().setStateName(valueStateName)
+ val message = ValueStateCall.newBuilder().setStateName(stateName)
.setClear(Clear.newBuilder().build()).build()
stateServer.handleValueStateRequest(message)
verify(valueState).clear()
verify(outputStream).writeInt(0)
}
test("value state update") {
- // Below byte array is a serialized row with a single integer value 1.
- val byteArray: Array[Byte] = Array(0x80.toByte, 0x05.toByte, 0x95.toByte, 0x05.toByte,
- 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte,
- 'K'.toByte, 0x01.toByte, 0x85.toByte, 0x94.toByte, '.'.toByte
- )
val byteString: ByteString = ByteString.copyFrom(byteArray)
- val message = ValueStateCall.newBuilder().setStateName(valueStateName)
+ val message = ValueStateCall.newBuilder().setStateName(stateName)
.setValueStateUpdate(ValueStateUpdate.newBuilder().setValue(byteString).build()).build()
stateServer.handleValueStateRequest(message)
verify(valueState).update(any[Row])
verify(outputStream).writeInt(0)
}
+
+ test("list state exists") {
+ val message = ListStateCall.newBuilder().setStateName(stateName)
+ .setExists(Exists.newBuilder().build()).build()
+ stateServer.handleListStateRequest(message)
+ verify(listState).exists()
+ }
+
+ test("list state get - iterator in map") {
+ val message = ListStateCall.newBuilder().setStateName(stateName)
+ .setListStateGet(ListStateGet.newBuilder().setIteratorId(iteratorId).build()).build()
+ stateServer.handleListStateRequest(message)
+ verify(listState, times(0)).get()
+ verify(arrowStreamWriter).writeRow(any)
+ verify(arrowStreamWriter).finalizeCurrentArrowBatch()
+ }
+
+ test("list state get - iterator not in map") {
+ val maxRecordsPerBatch = 2
+ val message = ListStateCall.newBuilder().setStateName(stateName)
+ .setListStateGet(ListStateGet.newBuilder().setIteratorId(iteratorId).build()).build()
+ val iteratorMap: mutable.HashMap[String, Iterator[Row]] = mutable.HashMap()
+ stateServer = new TransformWithStateInPandasStateServer(serverSocket,
+ statefulProcessorHandle, groupingKeySchema, "", false, false,
+ maxRecordsPerBatch, outputStream, valueStateMap,
+ transformWithStateInPandasDeserializer, arrowStreamWriter, listStateMap, iteratorMap)
+ when(listState.get()).thenReturn(Iterator(new GenericRowWithSchema(Array(1), stateSchema),
+ new GenericRowWithSchema(Array(2), stateSchema),
+ new GenericRowWithSchema(Array(3), stateSchema)))
+ stateServer.handleListStateRequest(message)
+ verify(listState).get()
+ // Verify that only maxRecordsPerBatch (2) rows are written to the output stream while still
+ // having 1 row left in the iterator.
+ verify(arrowStreamWriter, times(maxRecordsPerBatch)).writeRow(any)
+ verify(arrowStreamWriter).finalizeCurrentArrowBatch()
Review Comment:
Would we like to also test that the remaining part will be presented if we request the same iterator ID? Probably good to have a separate test to fully consume the iterator which length is multiple of maxRecordsPerBatch.
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
Review Comment:
nit: It's not for single value state. Probably could be an issue of copy & paste without fix.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -157,7 +185,11 @@ class TransformWithStateInPandasStateServer(
val ttlDurationMs = if (message.getGetValueState.hasTtl) {
Some(message.getGetValueState.getTtl.getDurationMs)
} else None
- initializeValueState(stateName, schema, ttlDurationMs)
+ initializeStateVariable(stateName, schema, "valueState", ttlDurationMs)
+ case StatefulProcessorCall.MethodCase.GETLISTSTATE =>
+ val stateName = message.getGetListState.getStateName
+ val schema = message.getGetListState.getSchema
+ initializeStateVariable(stateName, schema, "listState", None)
Review Comment:
Is it TODO to support TTL for ListState ? Probably better to have a JIRA ticket.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -232,23 +349,53 @@ class TransformWithStateInPandasStateServer(
outputStream.write(responseMessageBytes)
}
- private def initializeValueState(
- stateName: String,
- schemaString: String,
- ttlDurationMs: Option[Int]): Unit = {
- if (!valueStates.contains(stateName)) {
- val schema = StructType.fromString(schemaString)
- val state = if (ttlDurationMs.isEmpty) {
- statefulProcessorHandle.getValueState[Row](stateName, Encoders.row(schema))
- } else {
- statefulProcessorHandle.getValueState(
- stateName, Encoders.row(schema), TTLConfig(Duration.ofMillis(ttlDurationMs.get)))
- }
- val valueRowDeserializer = ExpressionEncoder(schema).resolveAndBind().createDeserializer()
- valueStates.put(stateName, (state, schema, valueRowDeserializer))
- sendResponse(0)
- } else {
- sendResponse(1, s"state $stateName already exists")
+ private def initializeStateVariable(
+ stateName: String,
+ schemaString: String,
+ stateVariable: String,
Review Comment:
nit: state`Type`? And would you mind using enum for type?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -201,13 +236,95 @@ class TransformWithStateInPandasStateServer(
}
case ValueStateCall.MethodCase.VALUESTATEUPDATE =>
val byteArray = message.getValueStateUpdate.getValue.toByteArray
- val valueStateTuple = valueStates(stateName)
// The value row is serialized as a byte array, we need to convert it back to a Row
- val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateTuple._2, valueStateTuple._3)
- valueStateTuple._1.update(valueRow)
+ val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateInfo.schema,
+ valueStateInfo.deserializer)
+ valueStateInfo.valueState.update(valueRow)
sendResponse(0)
case ValueStateCall.MethodCase.CLEAR =>
- valueStates(stateName)._1.clear()
+ valueStateInfo.valueState.clear()
+ sendResponse(0)
+ case _ =>
+ throw new IllegalArgumentException("Invalid method call")
+ }
+ }
+
+ private[sql] def handleListStateRequest(message: ListStateCall): Unit = {
+ val stateName = message.getStateName
+ if (!listStates.contains(stateName)) {
+ logWarning(log"List state ${MDC(LogKeys.STATE_NAME, stateName)} is not initialized.")
+ sendResponse(1, s"List state $stateName is not initialized.")
+ return
+ }
+ val listStateInfo = listStates(stateName)
+ val deserializer = if (deserializerForTest != null) {
+ deserializerForTest
+ } else {
+ new TransformWithStateInPandasDeserializer(listStateInfo.deserializer)
+ }
+ message.getMethodCase match {
+ case ListStateCall.MethodCase.EXISTS =>
+ if (listStateInfo.listState.exists()) {
+ sendResponse(0)
+ } else {
+ // Send status code 2 to indicate that the list state doesn't have a value yet.
+ sendResponse(2, s"state $stateName doesn't exist")
+ }
+ case ListStateCall.MethodCase.LISTSTATEPUT =>
+ val rows = deserializer.readArrowBatches(inputStream)
+ listStateInfo.listState.put(rows.toArray)
+ sendResponse(0)
+ case ListStateCall.MethodCase.LISTSTATEGET =>
+ val iteratorId = message.getListStateGet.getIteratorId
+ var iteratorOption = listStateIterators.get(iteratorId)
+ if (iteratorOption.isEmpty) {
+ iteratorOption = Some(listStateInfo.listState.get())
+ listStateIterators.put(iteratorId, iteratorOption.get)
+ }
+ if (!iteratorOption.get.hasNext) {
+ sendResponse(2, s"List state $stateName doesn't contain any value.")
+ return
+ } else {
+ sendResponse(0)
+ }
+ outputStream.flush()
+ val arrowStreamWriter = if (arrowStreamWriterForTest != null) {
+ arrowStreamWriterForTest
+ } else {
+ val arrowSchema = ArrowUtils.toArrowSchema(listStateInfo.schema, timeZoneId,
+ errorOnDuplicatedFieldNames, largeVarTypes)
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator(
+ s"stdout writer for transformWithStateInPandas state socket", 0, Long.MaxValue)
+ val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ new BaseStreamingArrowWriter(root, new ArrowStreamWriter(root, null, outputStream),
+ arrowTransformWithStateInPandasMaxRecordsPerBatch)
+ }
+ val listRowSerializer = listStateInfo.serializer
+ // Only write a single batch in each GET request. Stops writing row if rowCount reaches
+ // the arrowTransformWithStateInPandasMaxRecordsPerBatch limit. This is to avoid a case
Review Comment:
Is this to avoid the case, or to support the case? From the code implementation I feel like it's to support the case. Otherwise do we have the logic to block such a request?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -232,23 +349,53 @@ class TransformWithStateInPandasStateServer(
outputStream.write(responseMessageBytes)
}
- private def initializeValueState(
- stateName: String,
- schemaString: String,
- ttlDurationMs: Option[Int]): Unit = {
- if (!valueStates.contains(stateName)) {
- val schema = StructType.fromString(schemaString)
- val state = if (ttlDurationMs.isEmpty) {
- statefulProcessorHandle.getValueState[Row](stateName, Encoders.row(schema))
- } else {
- statefulProcessorHandle.getValueState(
- stateName, Encoders.row(schema), TTLConfig(Duration.ofMillis(ttlDurationMs.get)))
- }
- val valueRowDeserializer = ExpressionEncoder(schema).resolveAndBind().createDeserializer()
- valueStates.put(stateName, (state, schema, valueRowDeserializer))
- sendResponse(0)
- } else {
- sendResponse(1, s"state $stateName already exists")
+ private def initializeStateVariable(
+ stateName: String,
Review Comment:
nit: two more spaces for params
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -232,23 +349,53 @@ class TransformWithStateInPandasStateServer(
outputStream.write(responseMessageBytes)
}
- private def initializeValueState(
- stateName: String,
- schemaString: String,
- ttlDurationMs: Option[Int]): Unit = {
- if (!valueStates.contains(stateName)) {
- val schema = StructType.fromString(schemaString)
- val state = if (ttlDurationMs.isEmpty) {
- statefulProcessorHandle.getValueState[Row](stateName, Encoders.row(schema))
- } else {
- statefulProcessorHandle.getValueState(
- stateName, Encoders.row(schema), TTLConfig(Duration.ofMillis(ttlDurationMs.get)))
- }
- val valueRowDeserializer = ExpressionEncoder(schema).resolveAndBind().createDeserializer()
- valueStates.put(stateName, (state, schema, valueRowDeserializer))
- sendResponse(0)
- } else {
- sendResponse(1, s"state $stateName already exists")
+ private def initializeStateVariable(
+ stateName: String,
+ schemaString: String,
+ stateVariable: String,
+ ttlDurationMs: Option[Int]): Unit = {
+ val schema = StructType.fromString(schemaString)
+ val expressionEncoder = ExpressionEncoder(schema).resolveAndBind()
+ stateVariable match {
+ case "valueState" => if (!valueStates.contains(stateName)) {
+ val state = if (ttlDurationMs.isEmpty) {
+ statefulProcessorHandle.getValueState[Row](stateName, Encoders.row(schema))
+ } else {
+ statefulProcessorHandle.getValueState(
+ stateName, Encoders.row(schema), TTLConfig(Duration.ofMillis(ttlDurationMs.get)))
+ }
+ valueStates.put(stateName,
+ ValueStateInfo(state, schema, expressionEncoder.createDeserializer()))
+ sendResponse(0)
+ } else {
+ sendResponse(1, s"Value state $stateName already exists")
+ }
+ case "listState" => if (!listStates.contains(stateName)) {
Review Comment:
nit: let's either assert ttlDurationMs is empty or leave code comment for TODO
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1775664739
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -99,25 +99,25 @@ def exists(self) -> bool:
"""
return self._list_state_client.exists(self._state_name)
- def get(self) -> Iterator[Row]:
+ def get(self) -> Iterator[Tuple]:
Review Comment:
Yep, this is what we are doing today, will apply to all state types in the future as well, thanks for the reminder!
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1769253275
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
Review Comment:
Good catch, updated, thanks!
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator[Row]:
+ """
+ Get list state with an iterator.
+ """
+ return ListStateIterator(self._list_state_client, self._state_name)
+
+ def put(self, new_state: List[Any]) -> None:
Review Comment:
Makes sense. Should we use `Tuple` for all the cases here? I see we are doing so in `GroupState` https://github.com/apache/spark/blob/f76a9b1135e748649bdb9a2104360f0dc533cc1f/python/pyspark/sql/streaming/state.py#L95
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -201,13 +236,95 @@ class TransformWithStateInPandasStateServer(
}
case ValueStateCall.MethodCase.VALUESTATEUPDATE =>
val byteArray = message.getValueStateUpdate.getValue.toByteArray
- val valueStateTuple = valueStates(stateName)
// The value row is serialized as a byte array, we need to convert it back to a Row
- val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateTuple._2, valueStateTuple._3)
- valueStateTuple._1.update(valueRow)
+ val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateInfo.schema,
+ valueStateInfo.deserializer)
+ valueStateInfo.valueState.update(valueRow)
sendResponse(0)
case ValueStateCall.MethodCase.CLEAR =>
- valueStates(stateName)._1.clear()
+ valueStateInfo.valueState.clear()
+ sendResponse(0)
+ case _ =>
+ throw new IllegalArgumentException("Invalid method call")
+ }
+ }
+
+ private[sql] def handleListStateRequest(message: ListStateCall): Unit = {
+ val stateName = message.getStateName
+ if (!listStates.contains(stateName)) {
+ logWarning(log"List state ${MDC(LogKeys.STATE_NAME, stateName)} is not initialized.")
+ sendResponse(1, s"List state $stateName is not initialized.")
+ return
+ }
+ val listStateInfo = listStates(stateName)
+ val deserializer = if (deserializerForTest != null) {
+ deserializerForTest
+ } else {
+ new TransformWithStateInPandasDeserializer(listStateInfo.deserializer)
+ }
+ message.getMethodCase match {
+ case ListStateCall.MethodCase.EXISTS =>
+ if (listStateInfo.listState.exists()) {
+ sendResponse(0)
+ } else {
+ // Send status code 2 to indicate that the list state doesn't have a value yet.
+ sendResponse(2, s"state $stateName doesn't exist")
+ }
+ case ListStateCall.MethodCase.LISTSTATEPUT =>
+ val rows = deserializer.readArrowBatches(inputStream)
+ listStateInfo.listState.put(rows.toArray)
+ sendResponse(0)
+ case ListStateCall.MethodCase.LISTSTATEGET =>
+ val iteratorId = message.getListStateGet.getIteratorId
+ var iteratorOption = listStateIterators.get(iteratorId)
+ if (iteratorOption.isEmpty) {
+ iteratorOption = Some(listStateInfo.listState.get())
+ listStateIterators.put(iteratorId, iteratorOption.get)
+ }
+ if (!iteratorOption.get.hasNext) {
+ sendResponse(2, s"List state $stateName doesn't contain any value.")
+ return
+ } else {
+ sendResponse(0)
+ }
+ outputStream.flush()
+ val arrowStreamWriter = if (arrowStreamWriterForTest != null) {
+ arrowStreamWriterForTest
+ } else {
+ val arrowSchema = ArrowUtils.toArrowSchema(listStateInfo.schema, timeZoneId,
+ errorOnDuplicatedFieldNames, largeVarTypes)
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator(
+ s"stdout writer for transformWithStateInPandas state socket", 0, Long.MaxValue)
+ val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ new BaseStreamingArrowWriter(root, new ArrowStreamWriter(root, null, outputStream),
+ arrowTransformWithStateInPandasMaxRecordsPerBatch)
+ }
+ val listRowSerializer = listStateInfo.serializer
+ // Only write a single batch in each GET request. Stops writing row if rowCount reaches
+ // the arrowTransformWithStateInPandasMaxRecordsPerBatch limit. This is to avoid a case
Review Comment:
Sorry about the misleading wording, I meant to handle/support such a case, updated.
##########
python/pyspark/sql/streaming/stateful_processor_api_client.py:
##########
@@ -168,3 +207,43 @@ def _serialize_to_bytes(self, schema: StructType, data: Tuple) -> bytes:
def _deserialize_from_bytes(self, value: bytes) -> Any:
return self.pickleSer.loads(value)
+
+ def _send_arrow_state(self, schema: StructType, state: List[Tuple]) -> None:
Review Comment:
We have something similar in `types.py`: `_from_numpy_type` https://github.com/apache/spark/blob/f76a9b1135e748649bdb9a2104360f0dc533cc1f/python/pyspark/sql/types.py#L2166 which is opposite to the use case in this PR. Should we move these utility functions there?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -157,7 +185,11 @@ class TransformWithStateInPandasStateServer(
val ttlDurationMs = if (message.getGetValueState.hasTtl) {
Some(message.getGetValueState.getTtl.getDurationMs)
} else None
- initializeValueState(stateName, schema, ttlDurationMs)
+ initializeStateVariable(stateName, schema, "valueState", ttlDurationMs)
+ case StatefulProcessorCall.MethodCase.GETLISTSTATE =>
+ val stateName = message.getGetListState.getStateName
+ val schema = message.getGetListState.getSchema
+ initializeStateVariable(stateName, schema, "listState", None)
Review Comment:
Added. Btw, I have the TTL change in my local as well, do you prefer to combine it together with the current PR? I think separate PRs makes more sense but let me know if you have other ideas.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -125,9 +149,13 @@ class TransformWithStateInPandasStateServer(
// The key row is serialized as a byte array, we need to convert it back to a Row
val keyRow = PythonSQLUtils.toJVMRow(keyBytes, groupingKeySchema, keyRowDeserializer)
ImplicitGroupingKeyTracker.setImplicitKey(keyRow)
+ // Reset the list state iterators for a new grouping key.
+ listStateIterators = new mutable.HashMap[String, Iterator[Row]]()
Review Comment:
In unit tests, we only access the iterator map when calling `handleListStateRequest` which doesn't contain logic to reset the map. I think as long as we don't explicitly combine `ImplicitGroupingKeyRequest` and `ListStateRequest` together in unit test, we should be good. I added a note in the 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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1772479914
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -99,25 +99,25 @@ def exists(self) -> bool:
"""
return self._list_state_client.exists(self._state_name)
- def get(self) -> Iterator[Row]:
+ def get(self) -> Iterator[Tuple]:
Review Comment:
I'm OK with Tuple for user facing. We would store this as Row in state store, right? We shouldn't rely on pickling.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1759635501
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -197,13 +232,92 @@ class TransformWithStateInPandasStateServer(
}
case ValueStateCall.MethodCase.VALUESTATEUPDATE =>
val byteArray = message.getValueStateUpdate.getValue.toByteArray
- val valueStateTuple = valueStates(stateName)
// The value row is serialized as a byte array, we need to convert it back to a Row
- val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateTuple._2, valueStateTuple._3)
- valueStateTuple._1.update(valueRow)
+ val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateInfo.schema,
+ valueStateInfo.deserializer)
+ valueStateInfo.valueState.update(valueRow)
sendResponse(0)
case ValueStateCall.MethodCase.CLEAR =>
- valueStates(stateName)._1.clear()
+ valueStateInfo.valueState.clear()
+ sendResponse(0)
+ case _ =>
+ throw new IllegalArgumentException("Invalid method call")
+ }
+ }
+
+ private[sql] def handleListStateRequest(message: ListStateCall): Unit = {
+ val stateName = message.getStateName
+ if (!listStates.contains(stateName)) {
+ logWarning(log"List state ${MDC(LogKeys.STATE_NAME, stateName)} is not initialized.")
+ sendResponse(1, s"List state $stateName is not initialized.")
+ return
+ }
+ val listStateInfo = listStates(stateName)
+ val deserializer = if (deserializerForTest != null) {
+ deserializerForTest
+ } else {
+ new TransformWithStateInPandasDeserializer(listStateInfo.deserializer)
+ }
+ message.getMethodCase match {
+ case ListStateCall.MethodCase.EXISTS =>
+ if (listStateInfo.listState.exists()) {
+ sendResponse(0)
+ } else {
+ // Send status code 2 to indicate that the list state doesn't have a value yet.
+ sendResponse(2, s"state $stateName doesn't exist")
+ }
+ case ListStateCall.MethodCase.LISTSTATEPUT =>
+ val rows = deserializer.readArrowBatches(inputStream)
+ listStateInfo.listState.put(rows.toArray)
+ sendResponse(0)
+ case ListStateCall.MethodCase.GET =>
+ var iteratorOption = listStateIterators.get(stateName)
+ if (iteratorOption.isEmpty) {
+ iteratorOption = Some(listStateInfo.listState.get())
+ listStateIterators.put(stateName, iteratorOption.get)
+ }
+ if (!iteratorOption.get.hasNext) {
+ sendResponse(2, s"List state $stateName doesn't contain any value.")
+ return
+ } else {
+ sendResponse(0)
+ }
+ outputStream.flush()
+ val arrowStreamWriter = if (arrowStreamWriterForTest != null) {
+ arrowStreamWriterForTest
+ } else {
+ val arrowSchema = ArrowUtils.toArrowSchema(listStateInfo.schema, timeZoneId,
+ errorOnDuplicatedFieldNames, largeVarTypes)
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator(
+ s"stdout writer for transformWithStateInPandas state socket", 0, Long.MaxValue)
+ val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ new BaseStreamingArrowWriter(root, new ArrowStreamWriter(root, null, outputStream),
+ arrowTransformWithStateInPandasMaxRecordsPerBatch)
+ }
+ val listRowSerializer = listStateInfo.serializer
+ // Only write a single batch in each GET request. Stops writing row if rowCount reaches
Review Comment:
Updated, 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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746093738
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -62,6 +74,22 @@ class TransformWithStateInPandasStateServer(
new mutable.HashMap[String, (ValueState[Row], StructType,
ExpressionEncoder.Deserializer[Row])]()
}
+ // A map to store the list state name -> (list state, schema, list state row deserializer,
+ // list state row serializer) mapping.
+ private val listStates = if (listStatesMapForTest != null) {
+ listStatesMapForTest
+ } else {
+ new mutable.HashMap[String, (ListState[Row], StructType,
+ ExpressionEncoder.Deserializer[Row], ExpressionEncoder.Serializer[Row])]()
+ }
+ // A map to store the list state name -> iterator mapping. This is to keep track of the
+ // current iterator position for each list state in a grouping key in case user tries to fetch
+ // another list state before the current iterator is exhausted.
+ private var listStateIterators = if (listStateIteratorMapForTest != null) {
+ listStateIteratorMapForTest
+ } else {
+ new mutable.HashMap[String, Iterator[Row]]()
Review Comment:
I am also thinking out loud here by trying to reason about the life cycle of maintaining this map if we restart. An instance of `TransformWithStateInPandasStateServer` will be alive during the lifecycle of one call of UDF/batch, and we'll be accessing the iterator inside one call of `handleInputRows`. So we should be fine for the below two scenarios:
- Two different calls of `listState.get()` on the same grouping key, and they will be in two different batches. We expect the iterator to start from the beginning. As for different batches, the server instances are different, so we'll get two iterators starting from the beginning and this is expected.
- Query restart. Similar as above, if it is different batch, the iterator will always start from the beginning and this is also expected.
So this implementation of maintaining a local map of iterator should be acting as expected.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746355250
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator["PandasDataFrameLike"]:
Review Comment:
cc. @anishshri-db
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746370362
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator["PandasDataFrameLike"]:
Review Comment:
I circled back with `Iterator["PandasDataFrameLike"]` mainly because we provide the same parameter in `handleInputRows` where users can directly use `rows: Iterator["PandasDataFrameLike"]` to construct list state for `put()` and `appendList()`.
```
def handleInputRows(
self, key: Any, rows: Iterator["PandasDataFrameLike"]
) -> Iterator["PandasDataFrameLike"]:
```
If we only support `Iterator[Row]` for `put()` and `appendList()`, users might need extra transformation to achieve the same. @anishshri-db do you know if that's a common use case?
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1770700474
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -157,7 +185,11 @@ class TransformWithStateInPandasStateServer(
val ttlDurationMs = if (message.getGetValueState.hasTtl) {
Some(message.getGetValueState.getTtl.getDurationMs)
} else None
- initializeValueState(stateName, schema, ttlDurationMs)
+ initializeStateVariable(stateName, schema, "valueState", ttlDurationMs)
+ case StatefulProcessorCall.MethodCase.GETLISTSTATE =>
+ val stateName = message.getGetListState.getStateName
+ val schema = message.getGetListState.getSchema
+ initializeStateVariable(stateName, schema, "listState", None)
Review Comment:
I'm OK with having another PR. Let's just make sure we don't miss.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on PR #47933:
URL: https://github.com/apache/spark/pull/47933#issuecomment-2350763921
@HeartSaVioR @HyukjinKwon Could you help review this PR? 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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1737296085
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator["PandasDataFrameLike"]:
Review Comment:
@HeartSaVioR I'm keeping the `Iterator["PandasDataFrameLike"]` for `get()` for now since I feel it's easier for users to use `PandasDataFrameLike` as input parameter for `put()` and `appendList()`. It's better to keep it consistent for `get()` as well. But let me know if you have different ideas, 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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on PR #47933:
URL: https://github.com/apache/spark/pull/47933#issuecomment-2372407493
> Still linter failure happened. https://github.com/bogao007/spark/actions/runs/11004183619/job/30554657636
Fixed.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1774303338
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -50,19 +50,11 @@ def exists(self) -> bool:
"""
return self._value_state_client.exists(self._state_name)
- def get(self) -> Optional[Row]:
+ def get(self) -> Optional[Tuple]:
"""
Get the state value if it exists. Returns None if the state variable does not have a value.
"""
- value = self._value_state_client.get(self._state_name)
- if value is None:
- return None
- schema = self.schema
- if isinstance(schema, str):
- schema = cast(StructType, _parse_datatype_string(schema))
- # Create the Row using the values and schema fields
- row = _create_row(schema.fieldNames(), value)
- return row
+ return self._value_state_client.get(self._state_name)
def update(self, new_value: Any) -> None:
Review Comment:
nit: mind updating this to Tuple as well? Please address this as follow-up PR as this PR has been around for a long time and it's just a tiny 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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR closed pull request #47933: [SPARK-49463] Support ListState for TransformWithStateInPandas
URL: https://github.com/apache/spark/pull/47933
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746097930
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -124,6 +152,8 @@ class TransformWithStateInPandasStateServer(
// The key row is serialized as a byte array, we need to convert it back to a Row
val keyRow = PythonSQLUtils.toJVMRow(keyBytes, groupingKeySchema, keyRowDeserializer)
ImplicitGroupingKeyTracker.setImplicitKey(keyRow)
+ // Reset the list state iterators for a new grouping key.
+ listStateIterators = new mutable.HashMap[String, Iterator[Row]]()
sendResponse(0)
case ImplicitGroupingKeyRequest.MethodCase.REMOVEIMPLICITKEY =>
ImplicitGroupingKeyTracker.removeImplicitKey()
Review Comment:
We should probably also do listStateIterators = new mutable.HashMap[String, Iterator[Row]]() here.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746512565
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator["PandasDataFrameLike"]:
Review Comment:
We consider Spark SQL's Row as the base type of the state. ValueState followed the underlying implementation, and I expect all types to follow the same.
I'm OK to accept PandasDataFrameLike, but it should be secondary type of support (should be `Union[Row, PandasDataFrameLike]`) and we need to take care of conversion to finally store them to "Row". (We need to error out when given pandas DataFrame does not fit with the state schema.)
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1759635751
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator["PandasDataFrameLike"]:
Review Comment:
Updated to Iterator[Row]
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1737296085
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator["PandasDataFrameLike"]:
Review Comment:
@HeartSaVioR I'm keeping the `Iterator["PandasDataFrameLike"]` for `get()` for now since I feel it's easier for users to use `PandasDataFrameLike` as a input parameter for `put()` and `appendList()`. It's better to keep it consistent for `get()` as well. But let me know if you have different ideas, 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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #47933:
URL: https://github.com/apache/spark/pull/47933#issuecomment-2369989083
Still linter failure happened.
https://github.com/bogao007/spark/actions/runs/11004183619/job/30554657636
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1772343998
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -99,25 +99,25 @@ def exists(self) -> bool:
"""
return self._list_state_client.exists(self._state_name)
- def get(self) -> Iterator[Row]:
+ def get(self) -> Iterator[Tuple]:
Review Comment:
I'm leaning towards using Tuple to be consistent with both applyInPandasWithState and input/output types. I've also updated ValueState to return Tuple for consistency. But let me know if you have any concern on this, thanks! cc @HeartSaVioR @HyukjinKwon
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1769332084
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -232,23 +349,53 @@ class TransformWithStateInPandasStateServer(
outputStream.write(responseMessageBytes)
}
- private def initializeValueState(
- stateName: String,
- schemaString: String,
- ttlDurationMs: Option[Int]): Unit = {
- if (!valueStates.contains(stateName)) {
- val schema = StructType.fromString(schemaString)
- val state = if (ttlDurationMs.isEmpty) {
- statefulProcessorHandle.getValueState[Row](stateName, Encoders.row(schema))
- } else {
- statefulProcessorHandle.getValueState(
- stateName, Encoders.row(schema), TTLConfig(Duration.ofMillis(ttlDurationMs.get)))
- }
- val valueRowDeserializer = ExpressionEncoder(schema).resolveAndBind().createDeserializer()
- valueStates.put(stateName, (state, schema, valueRowDeserializer))
- sendResponse(0)
- } else {
- sendResponse(1, s"state $stateName already exists")
+ private def initializeStateVariable(
+ stateName: String,
+ schemaString: String,
+ stateVariable: String,
Review Comment:
done
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -232,23 +349,53 @@ class TransformWithStateInPandasStateServer(
outputStream.write(responseMessageBytes)
}
- private def initializeValueState(
- stateName: String,
- schemaString: String,
- ttlDurationMs: Option[Int]): Unit = {
- if (!valueStates.contains(stateName)) {
- val schema = StructType.fromString(schemaString)
- val state = if (ttlDurationMs.isEmpty) {
- statefulProcessorHandle.getValueState[Row](stateName, Encoders.row(schema))
- } else {
- statefulProcessorHandle.getValueState(
- stateName, Encoders.row(schema), TTLConfig(Duration.ofMillis(ttlDurationMs.get)))
- }
- val valueRowDeserializer = ExpressionEncoder(schema).resolveAndBind().createDeserializer()
- valueStates.put(stateName, (state, schema, valueRowDeserializer))
- sendResponse(0)
- } else {
- sendResponse(1, s"state $stateName already exists")
+ private def initializeStateVariable(
+ stateName: String,
+ schemaString: String,
+ stateVariable: String,
+ ttlDurationMs: Option[Int]): Unit = {
+ val schema = StructType.fromString(schemaString)
+ val expressionEncoder = ExpressionEncoder(schema).resolveAndBind()
+ stateVariable match {
+ case "valueState" => if (!valueStates.contains(stateName)) {
+ val state = if (ttlDurationMs.isEmpty) {
+ statefulProcessorHandle.getValueState[Row](stateName, Encoders.row(schema))
+ } else {
+ statefulProcessorHandle.getValueState(
+ stateName, Encoders.row(schema), TTLConfig(Duration.ofMillis(ttlDurationMs.get)))
+ }
+ valueStates.put(stateName,
+ ValueStateInfo(state, schema, expressionEncoder.createDeserializer()))
+ sendResponse(0)
+ } else {
+ sendResponse(1, s"Value state $stateName already exists")
+ }
+ case "listState" => if (!listStates.contains(stateName)) {
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1773807929
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -99,25 +99,25 @@ def exists(self) -> bool:
"""
return self._list_state_client.exists(self._state_name)
- def get(self) -> Iterator[Row]:
+ def get(self) -> Iterator[Tuple]:
Review Comment:
@HeartSaVioR Yep, we store as `Row` in state store. The pickling thing is that I followed what ApplyInPandasWithStateWriter does for groupState. It serializes `Row` state to bytes in JVM: https://github.com/apache/spark/blob/55d0233d19cc52bee91a9619057d9b6f33165a0a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ApplyInPandasWithStateWriter.scala#L182 and then deserialize using pickling in python: https://github.com/apache/spark/blob/55d0233d19cc52bee91a9619057d9b6f33165a0a/python/pyspark/sql/pandas/serializers.py#L837 Are you saying we shouldn't use `pickleSer.loads` in valueState `get()`? If so, could you share your concerns here?
An alternative I can think about is to let valueState also use arrow to transmit the state row from JVM to python side and deserialize just like list state. Let me know your thoughts on this, 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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #47933:
URL: https://github.com/apache/spark/pull/47933#issuecomment-2372737331
Thanks! Merging 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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1755201977
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -197,13 +232,92 @@ class TransformWithStateInPandasStateServer(
}
case ValueStateCall.MethodCase.VALUESTATEUPDATE =>
val byteArray = message.getValueStateUpdate.getValue.toByteArray
- val valueStateTuple = valueStates(stateName)
// The value row is serialized as a byte array, we need to convert it back to a Row
- val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateTuple._2, valueStateTuple._3)
- valueStateTuple._1.update(valueRow)
+ val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateInfo.schema,
+ valueStateInfo.deserializer)
+ valueStateInfo.valueState.update(valueRow)
sendResponse(0)
case ValueStateCall.MethodCase.CLEAR =>
- valueStates(stateName)._1.clear()
+ valueStateInfo.valueState.clear()
+ sendResponse(0)
+ case _ =>
+ throw new IllegalArgumentException("Invalid method call")
+ }
+ }
+
+ private[sql] def handleListStateRequest(message: ListStateCall): Unit = {
+ val stateName = message.getStateName
+ if (!listStates.contains(stateName)) {
+ logWarning(log"List state ${MDC(LogKeys.STATE_NAME, stateName)} is not initialized.")
+ sendResponse(1, s"List state $stateName is not initialized.")
+ return
+ }
+ val listStateInfo = listStates(stateName)
+ val deserializer = if (deserializerForTest != null) {
+ deserializerForTest
+ } else {
+ new TransformWithStateInPandasDeserializer(listStateInfo.deserializer)
+ }
+ message.getMethodCase match {
+ case ListStateCall.MethodCase.EXISTS =>
+ if (listStateInfo.listState.exists()) {
+ sendResponse(0)
+ } else {
+ // Send status code 2 to indicate that the list state doesn't have a value yet.
+ sendResponse(2, s"state $stateName doesn't exist")
+ }
+ case ListStateCall.MethodCase.LISTSTATEPUT =>
+ val rows = deserializer.readArrowBatches(inputStream)
+ listStateInfo.listState.put(rows.toArray)
+ sendResponse(0)
+ case ListStateCall.MethodCase.GET =>
+ var iteratorOption = listStateIterators.get(stateName)
+ if (iteratorOption.isEmpty) {
+ iteratorOption = Some(listStateInfo.listState.get())
+ listStateIterators.put(stateName, iteratorOption.get)
+ }
+ if (!iteratorOption.get.hasNext) {
+ sendResponse(2, s"List state $stateName doesn't contain any value.")
+ return
+ } else {
+ sendResponse(0)
+ }
+ outputStream.flush()
+ val arrowStreamWriter = if (arrowStreamWriterForTest != null) {
+ arrowStreamWriterForTest
+ } else {
+ val arrowSchema = ArrowUtils.toArrowSchema(listStateInfo.schema, timeZoneId,
+ errorOnDuplicatedFieldNames, largeVarTypes)
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator(
+ s"stdout writer for transformWithStateInPandas state socket", 0, Long.MaxValue)
+ val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ new BaseStreamingArrowWriter(root, new ArrowStreamWriter(root, null, outputStream),
+ arrowTransformWithStateInPandasMaxRecordsPerBatch)
+ }
+ val listRowSerializer = listStateInfo.serializer
+ // Only write a single batch in each GET request. Stops writing row if rowCount reaches
Review Comment:
nits: If the number of row exceed the `rowCount`, IIUC, we will just send the following data in another batch? Shall we log out some info here when the number of info exceeds `rowCount`? Also is it possible to write some tests for the above scenario to make sure the newly added config takes effect?
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746097674
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -124,6 +152,8 @@ class TransformWithStateInPandasStateServer(
// The key row is serialized as a byte array, we need to convert it back to a Row
val keyRow = PythonSQLUtils.toJVMRow(keyBytes, groupingKeySchema, keyRowDeserializer)
ImplicitGroupingKeyTracker.setImplicitKey(keyRow)
+ // Reset the list state iterators for a new grouping key.
+ listStateIterators = new mutable.HashMap[String, Iterator[Row]]()
sendResponse(0)
case ImplicitGroupingKeyRequest.MethodCase.REMOVEIMPLICITKEY =>
ImplicitGroupingKeyTracker.removeImplicitKey()
Review Comment:
We should probably also do `listStateIterators = new mutable.HashMap[String, Iterator[Row]]()` here.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -124,6 +152,8 @@ class TransformWithStateInPandasStateServer(
// The key row is serialized as a byte array, we need to convert it back to a Row
val keyRow = PythonSQLUtils.toJVMRow(keyBytes, groupingKeySchema, keyRowDeserializer)
ImplicitGroupingKeyTracker.setImplicitKey(keyRow)
+ // Reset the list state iterators for a new grouping key.
+ listStateIterators = new mutable.HashMap[String, Iterator[Row]]()
sendResponse(0)
case ImplicitGroupingKeyRequest.MethodCase.REMOVEIMPLICITKEY =>
ImplicitGroupingKeyTracker.removeImplicitKey()
Review Comment:
We should probably also do `listStateIterators = new mutable.HashMap[String, Iterator[Row]]()` here.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746355911
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator["PandasDataFrameLike"]:
Review Comment:
We decided to go with Row for get() in ValueState, right? I'm not sure we want to give different UX for ValueState vs ListState. It's conceptually just `List[ValueState]`, not a totally different once.
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -77,6 +78,58 @@ def clear(self) -> None:
self._value_state_client.clear(self._state_name)
+class ListState:
+ """
+ Class used for arbitrary stateful operations with transformWithState to capture single value
+ state.
+
+ .. versionadded:: 4.0.0
+ """
+
+ def __init__(
+ self, list_state_client: ListStateClient, state_name: str, schema: Union[StructType, str]
+ ) -> None:
+ self._list_state_client = list_state_client
+ self._state_name = state_name
+ self.schema = schema
+
+ def exists(self) -> bool:
+ """
+ Whether list state exists or not.
+ """
+ return self._list_state_client.exists(self._state_name)
+
+ def get(self) -> Iterator["PandasDataFrameLike"]:
Review Comment:
We decided to go with Row for get() in ValueState, right? I'm not sure we want to give different UX for ValueState vs ListState. It's conceptually just `List[ValueState]`, not a totally different one.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1746072835
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -62,6 +74,22 @@ class TransformWithStateInPandasStateServer(
new mutable.HashMap[String, (ValueState[Row], StructType,
ExpressionEncoder.Deserializer[Row])]()
}
+ // A map to store the list state name -> (list state, schema, list state row deserializer,
+ // list state row serializer) mapping.
+ private val listStates = if (listStatesMapForTest != null) {
+ listStatesMapForTest
+ } else {
+ new mutable.HashMap[String, (ListState[Row], StructType,
+ ExpressionEncoder.Deserializer[Row], ExpressionEncoder.Serializer[Row])]()
+ }
+ // A map to store the list state name -> iterator mapping. This is to keep track of the
+ // current iterator position for each list state in a grouping key in case user tries to fetch
+ // another list state before the current iterator is exhausted.
+ private var listStateIterators = if (listStateIteratorMapForTest != null) {
+ listStateIteratorMapForTest
+ } else {
+ new mutable.HashMap[String, Iterator[Row]]()
Review Comment:
Just thinking out loud here, curious if we should consider using thread-safe map here. I remember we were using thread-safety map in `RocksDBStateStoreProvider` here: https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/RocksDBStateStoreProvider.scala#L462
Though at the top of my head I think we should be fine here for `TransformWithStateInPandasStateServer` as each state store should have their own server instance so there won't be any race condition.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1755201977
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala:
##########
@@ -197,13 +232,92 @@ class TransformWithStateInPandasStateServer(
}
case ValueStateCall.MethodCase.VALUESTATEUPDATE =>
val byteArray = message.getValueStateUpdate.getValue.toByteArray
- val valueStateTuple = valueStates(stateName)
// The value row is serialized as a byte array, we need to convert it back to a Row
- val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateTuple._2, valueStateTuple._3)
- valueStateTuple._1.update(valueRow)
+ val valueRow = PythonSQLUtils.toJVMRow(byteArray, valueStateInfo.schema,
+ valueStateInfo.deserializer)
+ valueStateInfo.valueState.update(valueRow)
sendResponse(0)
case ValueStateCall.MethodCase.CLEAR =>
- valueStates(stateName)._1.clear()
+ valueStateInfo.valueState.clear()
+ sendResponse(0)
+ case _ =>
+ throw new IllegalArgumentException("Invalid method call")
+ }
+ }
+
+ private[sql] def handleListStateRequest(message: ListStateCall): Unit = {
+ val stateName = message.getStateName
+ if (!listStates.contains(stateName)) {
+ logWarning(log"List state ${MDC(LogKeys.STATE_NAME, stateName)} is not initialized.")
+ sendResponse(1, s"List state $stateName is not initialized.")
+ return
+ }
+ val listStateInfo = listStates(stateName)
+ val deserializer = if (deserializerForTest != null) {
+ deserializerForTest
+ } else {
+ new TransformWithStateInPandasDeserializer(listStateInfo.deserializer)
+ }
+ message.getMethodCase match {
+ case ListStateCall.MethodCase.EXISTS =>
+ if (listStateInfo.listState.exists()) {
+ sendResponse(0)
+ } else {
+ // Send status code 2 to indicate that the list state doesn't have a value yet.
+ sendResponse(2, s"state $stateName doesn't exist")
+ }
+ case ListStateCall.MethodCase.LISTSTATEPUT =>
+ val rows = deserializer.readArrowBatches(inputStream)
+ listStateInfo.listState.put(rows.toArray)
+ sendResponse(0)
+ case ListStateCall.MethodCase.GET =>
+ var iteratorOption = listStateIterators.get(stateName)
+ if (iteratorOption.isEmpty) {
+ iteratorOption = Some(listStateInfo.listState.get())
+ listStateIterators.put(stateName, iteratorOption.get)
+ }
+ if (!iteratorOption.get.hasNext) {
+ sendResponse(2, s"List state $stateName doesn't contain any value.")
+ return
+ } else {
+ sendResponse(0)
+ }
+ outputStream.flush()
+ val arrowStreamWriter = if (arrowStreamWriterForTest != null) {
+ arrowStreamWriterForTest
+ } else {
+ val arrowSchema = ArrowUtils.toArrowSchema(listStateInfo.schema, timeZoneId,
+ errorOnDuplicatedFieldNames, largeVarTypes)
+ val allocator = ArrowUtils.rootAllocator.newChildAllocator(
+ s"stdout writer for transformWithStateInPandas state socket", 0, Long.MaxValue)
+ val root = VectorSchemaRoot.create(arrowSchema, allocator)
+ new BaseStreamingArrowWriter(root, new ArrowStreamWriter(root, null, outputStream),
+ arrowTransformWithStateInPandasMaxRecordsPerBatch)
+ }
+ val listRowSerializer = listStateInfo.serializer
+ // Only write a single batch in each GET request. Stops writing row if rowCount reaches
Review Comment:
nits: If the number of row exceed the `rowCount`, IIUC, we will just send the following data in another batch? Shall we log out some info here when the number of rows exceeds `rowCount`? Also is it possible to write some tests for the above scenario to make sure the newly added config takes effect?
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "bogao007 (via GitHub)" <gi...@apache.org>.
bogao007 commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1769332212
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServerSuite.scala:
##########
@@ -127,24 +151,79 @@ class TransformWithStateInPandasStateServerSuite extends SparkFunSuite with Befo
}
test("value state clear") {
- val message = ValueStateCall.newBuilder().setStateName(valueStateName)
+ val message = ValueStateCall.newBuilder().setStateName(stateName)
.setClear(Clear.newBuilder().build()).build()
stateServer.handleValueStateRequest(message)
verify(valueState).clear()
verify(outputStream).writeInt(0)
}
test("value state update") {
- // Below byte array is a serialized row with a single integer value 1.
- val byteArray: Array[Byte] = Array(0x80.toByte, 0x05.toByte, 0x95.toByte, 0x05.toByte,
- 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte, 0x00.toByte,
- 'K'.toByte, 0x01.toByte, 0x85.toByte, 0x94.toByte, '.'.toByte
- )
val byteString: ByteString = ByteString.copyFrom(byteArray)
- val message = ValueStateCall.newBuilder().setStateName(valueStateName)
+ val message = ValueStateCall.newBuilder().setStateName(stateName)
.setValueStateUpdate(ValueStateUpdate.newBuilder().setValue(byteString).build()).build()
stateServer.handleValueStateRequest(message)
verify(valueState).update(any[Row])
verify(outputStream).writeInt(0)
}
+
+ test("list state exists") {
+ val message = ListStateCall.newBuilder().setStateName(stateName)
+ .setExists(Exists.newBuilder().build()).build()
+ stateServer.handleListStateRequest(message)
+ verify(listState).exists()
+ }
+
+ test("list state get - iterator in map") {
+ val message = ListStateCall.newBuilder().setStateName(stateName)
+ .setListStateGet(ListStateGet.newBuilder().setIteratorId(iteratorId).build()).build()
+ stateServer.handleListStateRequest(message)
+ verify(listState, times(0)).get()
+ verify(arrowStreamWriter).writeRow(any)
+ verify(arrowStreamWriter).finalizeCurrentArrowBatch()
+ }
+
+ test("list state get - iterator not in map") {
+ val maxRecordsPerBatch = 2
+ val message = ListStateCall.newBuilder().setStateName(stateName)
+ .setListStateGet(ListStateGet.newBuilder().setIteratorId(iteratorId).build()).build()
+ val iteratorMap: mutable.HashMap[String, Iterator[Row]] = mutable.HashMap()
+ stateServer = new TransformWithStateInPandasStateServer(serverSocket,
+ statefulProcessorHandle, groupingKeySchema, "", false, false,
+ maxRecordsPerBatch, outputStream, valueStateMap,
+ transformWithStateInPandasDeserializer, arrowStreamWriter, listStateMap, iteratorMap)
+ when(listState.get()).thenReturn(Iterator(new GenericRowWithSchema(Array(1), stateSchema),
+ new GenericRowWithSchema(Array(2), stateSchema),
+ new GenericRowWithSchema(Array(3), stateSchema)))
+ stateServer.handleListStateRequest(message)
+ verify(listState).get()
+ // Verify that only maxRecordsPerBatch (2) rows are written to the output stream while still
+ // having 1 row left in the iterator.
+ verify(arrowStreamWriter, times(maxRecordsPerBatch)).writeRow(any)
+ verify(arrowStreamWriter).finalizeCurrentArrowBatch()
Review Comment:
added a separate test case
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1769456350
##########
python/pyspark/sql/streaming/stateful_processor_api_client.py:
##########
@@ -168,3 +207,43 @@ def _serialize_to_bytes(self, schema: StructType, data: Tuple) -> bytes:
def _deserialize_from_bytes(self, value: bytes) -> Any:
return self.pickleSer.loads(value)
+
+ def _send_arrow_state(self, schema: StructType, state: List[Tuple]) -> None:
Review Comment:
Yeah let's move
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #47933:
URL: https://github.com/apache/spark/pull/47933#issuecomment-2367094990
https://github.com/bogao007/spark/actions/runs/10967170915/job/30456485161
linter error is a real one; looks like this came from Tuple/Row. We can address altogether.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1770702214
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -99,25 +99,25 @@ def exists(self) -> bool:
"""
return self._list_state_client.exists(self._state_name)
- def get(self) -> Iterator[Row]:
+ def get(self) -> Iterator[Tuple]:
Review Comment:
Wait, I see we provide "Row" in ValueState.get() and we build Row instance to match with the signature. Are we going to diverge UX, or the type hint is incorrect and list state also builds Row instance per element?
We don't need to strictly match with ApplyInPandasWithState but we need to be consistent among state types in TransformWithStateInPandas.
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -99,25 +99,25 @@ def exists(self) -> bool:
"""
return self._list_state_client.exists(self._state_name)
- def get(self) -> Iterator[Row]:
+ def get(self) -> Iterator[Tuple]:
Review Comment:
I'm OK either way, 1) be consistent with Tuple for state read & write 2) allow Tuple for state write but provide Row for state read to be strict with schema.
I don't even think it's a crazy idea to enforce Row for state write to be consistent. Though I kind of agree that Tuple is probably easier for users to deal with.
cc. @HyukjinKwon
Do we have preference in PySpark for such a scenario? We internally use Row, and need to decide whether we expose Row as it is, or whatever convenient type to users and handle conversion internally.
We previously used Tuple for applyInPandasWithState so probably the preference is the latter, but wanted to double confirm before moving on.
--
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
Re: [PR] [SPARK-49463] Support ListState for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #47933:
URL: https://github.com/apache/spark/pull/47933#discussion_r1774302216
##########
python/pyspark/sql/streaming/stateful_processor.py:
##########
@@ -99,25 +99,25 @@ def exists(self) -> bool:
"""
return self._list_state_client.exists(self._state_name)
- def get(self) -> Iterator[Row]:
+ def get(self) -> Iterator[Tuple]:
Review Comment:
No, what I meant was that even we get Tuple from users we should convert this to Row internally before sending to JVM, and vice versa. Let's ensure this would be applied for all state types. I'm going to assume this is true.
--
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