You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "jingz-db (via GitHub)" <gi...@apache.org> on 2024/12/11 23:56:22 UTC
[PR] [WIP] Python TWS metadata [spark]
jingz-db opened a new pull request, #49156:
URL: https://github.com/apache/spark/pull/49156
<!--
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?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
### 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'.
-->
### 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.
-->
### 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.
-->
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887742434
##########
python/pyspark/sql/streaming/transform_with_state_driver_worker.py:
##########
@@ -0,0 +1,89 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import json
+
+from pyspark.util import local_connect_and_auth
+from pyspark.serializers import (
+ write_int,
+ read_int,
+ UTF8Deserializer,
+ CPickleSerializer,
+)
+from pyspark import worker
+from pyspark.util import handle_worker_exception
+from typing import IO
+from pyspark.worker_util import check_python_version
+from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
+from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+from pyspark.sql.types import StructType
+
+pickle_ser = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def main(infile: IO, outfile: IO) -> None:
Review Comment:
Added a line below at line 44.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1894535573
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +920,359 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ # This test covers mapState with TTL, an empty state variable
+ # and additional test against initial state python runner
+ def test_transform_with_map_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
+ map_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "mapState")
+ .load()
+ )
+ assert set(
Review Comment:
nit: any reason to convert this as set while you sort to make the order deterministic?
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +920,359 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ # This test covers mapState with TTL, an empty state variable
+ # and additional test against initial state python runner
+ def test_transform_with_map_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
+ map_state_df = (
Review Comment:
Shall we check flatten type here as well, likewise we did for list state type?
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +920,359 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ # This test covers mapState with TTL, an empty state variable
+ # and additional test against initial state python runner
+ def test_transform_with_map_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
+ map_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "mapState")
+ .load()
+ )
+ assert set(
+ map_state_df.selectExpr(
+ "key.id AS groupingKey",
+ "user_map_key.name AS mapKey",
+ "user_map_value.value.count AS mapValue",
+ )
+ .sort("groupingKey")
+ .collect()
+ ) == {
+ Row(groupingKey="0", mapKey="key2", mapValue=2),
+ Row(groupingKey="1", mapKey="key2", mapValue=2),
+ }
+
+ ttl_df = map_state_df.selectExpr(
+ "user_map_value.ttlExpirationMs AS TTLVal"
+ ).collect()
+ # check if there are two rows containing TTL value in map state dataframe
+ assert len(ttl_df) == 2
+ # check if two rows are of the same TTL value
+ assert len(set(ttl_df)) == 1
+
+ list_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState")
+ .load()
+ )
+ assert list_state_df.isEmpty()
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # run the same test suite again but with no-op initial state
+ # TWS with initial state is using a different python runner
+ init_data = [("0", 789), ("3", 987)]
+ initial_state = self.spark.createDataFrame(init_data, "id string, temperature int").groupBy(
+ "id"
+ )
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
Review Comment:
I was wondering how the result could be same between non-initial state and initial state. So looks like we do not do anything with initial state in that processor, right? I'm OK with it, just to double confirm.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #49156:
URL: https://github.com/apache/spark/pull/49156#issuecomment-2606079406
Just a quick note .. Seems like the test `test_value_state_ttl_expiration` is still flaky when old dependencies are used (https://github.com/apache/spark/actions/runs/12883552117/job/35918143550).
@jingz-db mind taking a look please? I filed a JIRA at https://issues.apache.org/jira/browse/SPARK-50908
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887720226
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +224,130 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+class TransformWithStateInPandasPythonPreInitRunner(
+ func: PythonFunction,
+ workerModule: String,
+ timeZoneId: String,
+ groupingKeySchema: StructType,
+ processorHandleImpl: DriverStatefulProcessorHandleImpl)
+ extends StreamingPythonRunner(func, "", "", workerModule)
+ with TransformWithStateInPandasPythonRunnerUtils
+ with Logging {
+ protected val sqlConf = SQLConf.get
+
+ private var dataIn: DataInputStream = _
+ private var dataOut: DataOutputStream = _
+
+ private var daemonThread: Thread = _
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ val env = SparkEnv.get
+
+ val localdir = env.blockManager.diskBlockManager.localDirs.map(f => f.getPath()).mkString(",")
Review Comment:
Is this the preferred method to get the local dir ?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1896073496
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -358,7 +369,7 @@ def check_results(batch_df, batch_id):
for q in self.spark.streams.active:
q.stop()
if batch_id == 0 or batch_id == 1:
- time.sleep(6)
+ time.sleep(4)
Review Comment:
This is an intentional change because the CI test won't pass with prior 6 seconds. I assume this is because the pre-init stage is now added and it take a bit more time to finish the first batch so some TTL state already timed out before we get to check results on the second batch.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1902215696
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +223,108 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+/**
+ * TransformWithStateInPandas driver side Python runner. Similar as executor side runner,
+ * will start a new daemon thread on the Python runner to run state server.
+ */
+class TransformWithStateInPandasPythonPreInitRunner(
+ func: PythonFunction,
+ workerModule: String,
+ timeZoneId: String,
+ groupingKeySchema: StructType,
+ processorHandleImpl: DriverStatefulProcessorHandleImpl)
+ extends StreamingPythonRunner(func, "", "", workerModule)
+ with TransformWithStateInPandasPythonRunnerUtils
+ with Logging {
+ protected val sqlConf = SQLConf.get
+
+ private var dataOut: DataOutputStream = _
+ private var dataIn: DataInputStream = _
+
+ private var daemonThread: Thread = _
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ val result = super.init()
+ dataOut = result._1
+ dataIn = result._2
+
+ // start state server, update socket port
+ startStateServer()
+ (dataOut, dataIn)
+ }
+
+ def process(): Unit = {
+ // Also write the port number for state server
+ dataOut.writeInt(stateServerSocketPort)
+ PythonWorkerUtils.writeUTF(groupingKeySchema.json, dataOut)
+ dataOut.flush()
+
+ val resFromPython = dataIn.readInt()
+ if (resFromPython != 0) {
+ val errMessage = PythonWorkerUtils.readUTF(dataIn)
+ throw streamingPythonRunnerInitializationFailure(resFromPython, errMessage)
+ }
+ }
+
+ override def stop(): Unit = {
+ super.stop()
+ closeServerSocketChannelSilently(stateServerSocket)
+ daemonThread.stop()
Review Comment:
Thanks for the reminding! I'll resolve this and send out a followup PR next week.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887732419
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasExec.scala:
##########
@@ -91,6 +93,7 @@ case class TransformWithStateInPandasExec(
private val sessionLocalTimeZone = conf.sessionLocalTimeZone
private val pythonRunnerConf = ArrowPythonRunner.getPythonRunnerConfMap(conf)
private[this] val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
+ private val (dedupAttributes, argOffsets) = resolveArgOffsets(child.output, groupingAttributes)
Review Comment:
This is unchanged in the base change. It was originally here: https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasExec.scala#L318. I am only moving this out as a class variable so that other function will need to access it as well.
This function is mainly used to resolve the offsets for restoring key and value attributes in the python side of runner here: https://github.com/apache/spark/blob/master/python/pyspark/worker.py#L1706
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1889347155
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +923,184 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ def test_transform_with_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
Review Comment:
I've added more tests and leave some comments on what variation was covered in the test cases.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR closed pull request #49156: [SPARK-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas
URL: https://github.com/apache/spark/pull/49156
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1889358511
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +923,328 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ # This test covers mapState with TTL, an empty state variable
+ # and additional test against initial state python runner
+ def test_transform_with_map_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
+ map_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "mapState")
+ .load()
+ )
+ assert set(
+ map_state_df.selectExpr(
+ "key.id AS groupingKey",
+ "user_map_key.name AS mapKey",
+ "user_map_value.value.count AS mapValue",
+ )
+ .sort("groupingKey")
+ .collect()
+ ) == {
+ Row(groupingKey="0", mapKey="key2", mapValue=2),
+ Row(groupingKey="1", mapKey="key2", mapValue=2),
+ }
+
+ ttl_df = map_state_df.selectExpr(
+ "user_map_value.ttlExpirationMs AS TTLVal"
+ ).collect()
+ # check if there are two rows containing TTL value in map state dataframe
+ assert len(ttl_df) == 2
+ # check if two rows are of the same TTL value
+ assert len(set(ttl_df)) == 1
+
+ list_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState")
+ .load()
+ )
+ assert list_state_df.isEmpty()
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # run the same test suite again but with no-op initial state
+ # TWS with initial state is using a different python runner
+ init_data = [("0", 789), ("3", 987)]
+ initial_state = self.spark.createDataFrame(init_data, "id string, temperature int").groupBy(
+ "id"
+ )
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=initial_state,
+ )
+
+ # This test covers multiple list state variables and flatten option
+ def test_transform_with_list_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] in ["listState1", "listState2"]:
+ state_var["stateVariableType"] == "ListState"
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source and flatten option
+ list_state_1_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState1")
+ .option("flattenCollectionTypes", True)
+ .load()
+ )
+ assert list_state_1_df.selectExpr(
+ "key.id AS groupingKey",
+ "list_element.temperature AS listElement",
+ ).sort("groupingKey", "listElement").collect() == [
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=111),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=111),
+ Row(groupingKey="1", listElement=120),
+ Row(groupingKey="1", listElement=120),
+ ]
+
+ list_state_2_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState2")
+ .option("flattenCollectionTypes", False)
+ .load()
+ )
+ assert list_state_2_df.selectExpr(
+ "key.id AS groupingKey", "list_value.temperature AS valueList"
+ ).sort("groupingKey").withColumn(
+ "valueSortedList", array_sort(col("valueList"))
+ ).select(
+ "groupingKey", "valueSortedList"
+ ).collect() == [
+ Row(groupingKey="0", valueSortedList=[20, 20, 120, 120, 222]),
+ Row(groupingKey="1", valueSortedList=[20, 20, 120, 120, 222]),
+ ]
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ ListStateProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # This test covers value state variable and read change feed
+ def test_transform_with_value_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="3"),
+ Row(id="1", countAsString="2"),
+ }
+
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+
+ # TODO "tempState" should not appear in the metadata as it is already deleted
Review Comment:
Yea - we can take a look and make that a separate change, if needed
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887717851
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasExec.scala:
##########
@@ -91,6 +93,7 @@ case class TransformWithStateInPandasExec(
private val sessionLocalTimeZone = conf.sessionLocalTimeZone
private val pythonRunnerConf = ArrowPythonRunner.getPythonRunnerConfMap(conf)
private[this] val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
+ private val (dedupAttributes, argOffsets) = resolveArgOffsets(child.output, groupingAttributes)
Review Comment:
Hmm - what does this do internally ? could we add some comment maybe ?
##########
python/pyspark/sql/streaming/transform_with_state_driver_worker.py:
##########
@@ -0,0 +1,89 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import json
+
+from pyspark.util import local_connect_and_auth
+from pyspark.serializers import (
+ write_int,
+ read_int,
+ UTF8Deserializer,
+ CPickleSerializer,
+)
+from pyspark import worker
+from pyspark.util import handle_worker_exception
+from typing import IO
+from pyspark.worker_util import check_python_version
+from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
+from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+from pyspark.sql.types import StructType
+
+pickle_ser = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def main(infile: IO, outfile: IO) -> None:
+ check_python_version(infile)
+
+ log_name = "Streaming TransformWithStateInPandas Python worker"
+
+ def process(processor, mode, key, input):
+ print(f"{log_name} Started execution of UDF: {func}.\n")
+ func(processor, mode, key, input)
+ print(f"{log_name} Completed execution of UDF: {func}.\n")
+
+ try:
+ func, return_type = worker.read_command(pickle_ser, infile)
+ # send signal for getting args
+ write_int(0, outfile)
+ outfile.flush()
+
+ while True:
+ state_server_port = read_int(infile)
+ key_schema = StructType.fromJson(json.loads(utf8_deserializer.loads(infile)))
+ print(
+ f"{log_name} received parameters for UDF. State server port: {state_server_port}, "
+ f"key schema: {key_schema}.\n"
+ )
+
+ stateful_processor_api_client = StatefulProcessorApiClient(
+ state_server_port, key_schema
+ )
+ process(
+ stateful_processor_api_client,
+ TransformWithStateInPandasFuncMode.PRE_INIT,
+ None,
+ iter([]),
+ )
+ write_int(0, outfile)
+ outfile.flush()
+ except Exception as e:
+ handle_worker_exception(e, outfile)
+ outfile.flush()
+
+
+if __name__ == "__main__":
+ # Read information about how to connect back to the JVM from the environment.
+ java_port = int(os.environ["PYTHON_WORKER_FACTORY_PORT"])
+ auth_secret = os.environ["PYTHON_WORKER_FACTORY_SECRET"]
+ (sock_file, sock) = local_connect_and_auth(java_port, auth_secret)
+ # There could be a long time between each micro batch.
+ sock.settimeout(None)
Review Comment:
Will we wait indefinitely 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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887797306
##########
python/pyspark/sql/streaming/transform_with_state_driver_worker.py:
##########
@@ -0,0 +1,89 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import json
+
+from pyspark.util import local_connect_and_auth
+from pyspark.serializers import (
+ write_int,
+ read_int,
+ UTF8Deserializer,
+ CPickleSerializer,
+)
+from pyspark import worker
+from pyspark.util import handle_worker_exception
+from typing import IO
+from pyspark.worker_util import check_python_version
+from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
+from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+from pyspark.sql.types import StructType
+
+pickle_ser = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def main(infile: IO, outfile: IO) -> None:
+ check_python_version(infile)
+
+ log_name = "Streaming TransformWithStateInPandas Python worker"
+
+ def process(processor, mode, key, input):
+ print(f"{log_name} Started execution of UDF: {func}.\n")
+ func(processor, mode, key, input)
+ print(f"{log_name} Completed execution of UDF: {func}.\n")
+
+ try:
+ func, return_type = worker.read_command(pickle_ser, infile)
+ # send signal for getting args
+ write_int(0, outfile)
+ outfile.flush()
+
+ while True:
+ state_server_port = read_int(infile)
+ key_schema = StructType.fromJson(json.loads(utf8_deserializer.loads(infile)))
+ print(
+ f"{log_name} received parameters for UDF. State server port: {state_server_port}, "
+ f"key schema: {key_schema}.\n"
+ )
+
+ stateful_processor_api_client = StatefulProcessorApiClient(
+ state_server_port, key_schema
+ )
+ process(
+ stateful_processor_api_client,
+ TransformWithStateInPandasFuncMode.PRE_INIT,
+ None,
+ iter([]),
+ )
+ write_int(0, outfile)
+ outfile.flush()
+ except Exception as e:
+ handle_worker_exception(e, outfile)
+ outfile.flush()
+
+
+if __name__ == "__main__":
+ # Read information about how to connect back to the JVM from the environment.
+ java_port = int(os.environ["PYTHON_WORKER_FACTORY_PORT"])
+ auth_secret = os.environ["PYTHON_WORKER_FACTORY_SECRET"]
+ (sock_file, sock) = local_connect_and_auth(java_port, auth_secret)
+ # There could be a long time between each micro batch.
+ sock.settimeout(None)
Review Comment:
Maybe we can just follow the same pattern as the worker on the executor side ?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on PR #49156:
URL: https://github.com/apache/spark/pull/49156#issuecomment-2547147649
@jingz-db - test failures seem relevant ?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887719393
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +923,184 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ def test_transform_with_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
Review Comment:
Are we covering all variations here ?
- all supported composite types
- options supported by state data source reader
- change data feed related options
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887719393
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +923,184 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ def test_transform_with_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
Review Comment:
Are we covering all variations here ?
- all supported composite types
- options supported by state data source reader ?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887722002
##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -504,6 +504,20 @@ def transformWithStateInPandas(
if isinstance(outputStructType, str):
outputStructType = cast(StructType, _parse_datatype_string(outputStructType))
+ def handle_pre_init(
+ statefulProcessorApiClient: StatefulProcessorApiClient,
+ ) -> Iterator["PandasDataFrameLike"]:
+ # driver handle is different from the handle used on executors
+ driver_handle = StatefulProcessorHandle(statefulProcessorApiClient)
+ statefulProcessorApiClient.set_handle_state(StatefulProcessorHandleState.PRE_INIT)
+ statefulProcessor.init(driver_handle)
+
+ # TODO figure out why we need to call close here in Scala side
Review Comment:
Sounds good. I'll un-comment the statefulprocess.close() and add some comments 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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887742023
##########
python/pyspark/sql/streaming/stateful_processor_util.py:
##########
@@ -25,3 +25,4 @@ class TransformWithStateInPandasFuncMode(Enum):
PROCESS_DATA = 1
PROCESS_TIMER = 2
COMPLETE = 3
+ PRE_INIT = 4
Review Comment:
>could we add this as a comment in the code too ?
Yeah this is documented as class comment here: https://github.com/apache/spark/blob/master/python/pyspark/sql/streaming/stateful_processor_util.py#L20
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on PR #49156:
URL: https://github.com/apache/spark/pull/49156#issuecomment-2607941901
Sorry @HyukjinKwon, what is old dependencies you are referring to here? I think this suite is flaky because of timing issue. I'll take the ticket and fix the flakiness soon.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887797514
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +224,130 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+class TransformWithStateInPandasPythonPreInitRunner(
+ func: PythonFunction,
+ workerModule: String,
+ timeZoneId: String,
+ groupingKeySchema: StructType,
+ processorHandleImpl: DriverStatefulProcessorHandleImpl)
+ extends StreamingPythonRunner(func, "", "", workerModule)
+ with TransformWithStateInPandasPythonRunnerUtils
+ with Logging {
+ protected val sqlConf = SQLConf.get
+
+ private var dataIn: DataInputStream = _
+ private var dataOut: DataOutputStream = _
+
+ private var daemonThread: Thread = _
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ val env = SparkEnv.get
+
+ val localdir = env.blockManager.diskBlockManager.localDirs.map(f => f.getPath()).mkString(",")
Review Comment:
ok sounds good
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1889350550
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +923,328 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ # This test covers mapState with TTL, an empty state variable
+ # and additional test against initial state python runner
+ def test_transform_with_map_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
+ map_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "mapState")
+ .load()
+ )
+ assert set(
+ map_state_df.selectExpr(
+ "key.id AS groupingKey",
+ "user_map_key.name AS mapKey",
+ "user_map_value.value.count AS mapValue",
+ )
+ .sort("groupingKey")
+ .collect()
+ ) == {
+ Row(groupingKey="0", mapKey="key2", mapValue=2),
+ Row(groupingKey="1", mapKey="key2", mapValue=2),
+ }
+
+ ttl_df = map_state_df.selectExpr(
+ "user_map_value.ttlExpirationMs AS TTLVal"
+ ).collect()
+ # check if there are two rows containing TTL value in map state dataframe
+ assert len(ttl_df) == 2
+ # check if two rows are of the same TTL value
+ assert len(set(ttl_df)) == 1
+
+ list_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState")
+ .load()
+ )
+ assert list_state_df.isEmpty()
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # run the same test suite again but with no-op initial state
+ # TWS with initial state is using a different python runner
+ init_data = [("0", 789), ("3", 987)]
+ initial_state = self.spark.createDataFrame(init_data, "id string, temperature int").groupBy(
+ "id"
+ )
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=initial_state,
+ )
+
+ # This test covers multiple list state variables and flatten option
+ def test_transform_with_list_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] in ["listState1", "listState2"]:
+ state_var["stateVariableType"] == "ListState"
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source and flatten option
+ list_state_1_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState1")
+ .option("flattenCollectionTypes", True)
+ .load()
+ )
+ assert list_state_1_df.selectExpr(
+ "key.id AS groupingKey",
+ "list_element.temperature AS listElement",
+ ).sort("groupingKey", "listElement").collect() == [
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=111),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=111),
+ Row(groupingKey="1", listElement=120),
+ Row(groupingKey="1", listElement=120),
+ ]
+
+ list_state_2_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState2")
+ .option("flattenCollectionTypes", False)
+ .load()
+ )
+ assert list_state_2_df.selectExpr(
+ "key.id AS groupingKey", "list_value.temperature AS valueList"
+ ).sort("groupingKey").withColumn(
+ "valueSortedList", array_sort(col("valueList"))
+ ).select(
+ "groupingKey", "valueSortedList"
+ ).collect() == [
+ Row(groupingKey="0", valueSortedList=[20, 20, 120, 120, 222]),
+ Row(groupingKey="1", valueSortedList=[20, 20, 120, 120, 222]),
+ ]
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ ListStateProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # This test covers value state variable and read change feed
+ def test_transform_with_value_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="3"),
+ Row(id="1", countAsString="2"),
+ }
+
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+
+ # TODO "tempState" should not appear in the metadata as it is already deleted
Review Comment:
I personally feel this should not appear in the metadata if we call `deleteIfExists` on "tempState". Currently DriverStatefulProcessorHandleImpl did nothing when calling `deleteIfExists`: https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StatefulProcessorHandleImpl.scala#L493. Shall we make the change to remove the deleted variable from tracking map for metadata?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1896040613
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -852,13 +872,6 @@ def check_results(batch_df, batch_id):
StatefulProcessorWithInitialStateTimers(), check_results, "processingTime"
)
- # run the same test suites again but with single shuffle partition
Review Comment:
I am moving this to the end of the test suites for better organization of the file.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887722827
##########
python/pyspark/sql/streaming/stateful_processor_util.py:
##########
@@ -25,3 +25,4 @@ class TransformWithStateInPandasFuncMode(Enum):
PROCESS_DATA = 1
PROCESS_TIMER = 2
COMPLETE = 3
+ PRE_INIT = 4
Review Comment:
Ah ok - could we add this as a comment in the code too ?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1896075454
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +920,359 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ # This test covers mapState with TTL, an empty state variable
+ # and additional test against initial state python runner
+ def test_transform_with_map_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
+ map_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "mapState")
+ .load()
+ )
+ assert set(
+ map_state_df.selectExpr(
+ "key.id AS groupingKey",
+ "user_map_key.name AS mapKey",
+ "user_map_value.value.count AS mapValue",
+ )
+ .sort("groupingKey")
+ .collect()
+ ) == {
+ Row(groupingKey="0", mapKey="key2", mapValue=2),
+ Row(groupingKey="1", mapKey="key2", mapValue=2),
+ }
+
+ ttl_df = map_state_df.selectExpr(
+ "user_map_value.ttlExpirationMs AS TTLVal"
+ ).collect()
+ # check if there are two rows containing TTL value in map state dataframe
+ assert len(ttl_df) == 2
+ # check if two rows are of the same TTL value
+ assert len(set(ttl_df)) == 1
+
+ list_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState")
+ .load()
+ )
+ assert list_state_df.isEmpty()
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # run the same test suite again but with no-op initial state
+ # TWS with initial state is using a different python runner
+ init_data = [("0", 789), ("3", 987)]
+ initial_state = self.spark.createDataFrame(init_data, "id string, temperature int").groupBy(
+ "id"
+ )
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
Review Comment:
Yes, you are correct. I am adding this test explicitly for initial state because initial state is using a different python runner and it goes to a different code path in `group_ops.py`. So I want to double check this metadata is also working with initial state python runner. But the `MapStateLargeTTLProcessor` we use here for testing defines `handleInitialState()` as a no-op so the results will be the same as non-initial state.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #49156:
URL: https://github.com/apache/spark/pull/49156#issuecomment-2560594371
@jingz-db There seems to be test failure? https://github.com/jingz-db/spark/actions/runs/12474212623/job/34815951143
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1896075656
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +920,359 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ # This test covers mapState with TTL, an empty state variable
+ # and additional test against initial state python runner
+ def test_transform_with_map_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
+ map_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "mapState")
+ .load()
+ )
+ assert set(
+ map_state_df.selectExpr(
+ "key.id AS groupingKey",
+ "user_map_key.name AS mapKey",
+ "user_map_value.value.count AS mapValue",
+ )
+ .sort("groupingKey")
+ .collect()
+ ) == {
+ Row(groupingKey="0", mapKey="key2", mapValue=2),
+ Row(groupingKey="1", mapKey="key2", mapValue=2),
+ }
+
+ ttl_df = map_state_df.selectExpr(
+ "user_map_value.ttlExpirationMs AS TTLVal"
+ ).collect()
+ # check if there are two rows containing TTL value in map state dataframe
+ assert len(ttl_df) == 2
+ # check if two rows are of the same TTL value
+ assert len(set(ttl_df)) == 1
+
+ list_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState")
+ .load()
+ )
+ assert list_state_df.isEmpty()
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # run the same test suite again but with no-op initial state
+ # TWS with initial state is using a different python runner
+ init_data = [("0", 789), ("3", 987)]
+ initial_state = self.spark.createDataFrame(init_data, "id string, temperature int").groupBy(
+ "id"
+ )
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=initial_state,
+ )
+
+ # This test covers multiple list state variables and flatten option
+ def test_transform_with_list_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] in ["listState1", "listState2"]:
+ state_var["stateVariableType"] == "ListState"
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source and flatten option
+ list_state_1_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState1")
+ .option("flattenCollectionTypes", True)
+ .load()
+ )
+ assert list_state_1_df.selectExpr(
+ "key.id AS groupingKey",
+ "list_element.temperature AS listElement",
+ ).sort("groupingKey", "listElement").collect() == [
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=111),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=111),
+ Row(groupingKey="1", listElement=120),
+ Row(groupingKey="1", listElement=120),
+ ]
+
+ list_state_2_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState2")
+ .option("flattenCollectionTypes", False)
+ .load()
+ )
+ assert list_state_2_df.selectExpr(
+ "key.id AS groupingKey", "list_value.temperature AS valueList"
+ ).sort("groupingKey").withColumn(
+ "valueSortedList", array_sort(col("valueList"))
+ ).select(
+ "groupingKey", "valueSortedList"
+ ).collect() == [
+ Row(groupingKey="0", valueSortedList=[20, 20, 120, 120, 222]),
+ Row(groupingKey="1", valueSortedList=[20, 20, 120, 120, 222]),
+ ]
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ ListStateProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # This test covers value state variable and read change feed,
+ # snapshotStartBatchId related options
+ def test_transform_with_value_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="3"),
+ Row(id="1", countAsString="2"),
+ }
+
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] in ["numViolations", "tempState"]:
+ state_var["stateVariableType"] == "ValueState"
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source and readChangeFeed
+ value_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "numViolations")
+ .option("readChangeFeed", True)
+ .option("changeStartBatchId", 0)
+ .load()
+ ).selectExpr(
+ "change_type", "key.id AS groupingKey", "value.value AS value", "partition_id"
+ )
+
+ assert value_state_df.select("change_type", "groupingKey", "value").sort(
+ "groupingKey"
+ ).collect() == [
+ Row(change_type="update", groupingKey="0", value=1),
+ Row(change_type="update", groupingKey="1", value=2),
+ ]
+
+ partition_id_list = [
+ row["partition_id"] for row in value_state_df.select("partition_id").collect()
+ ]
+
+ for partition_id in partition_id_list:
+ # check for state data source and snapshotStartBatchId options
+ state_snapshot_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "numViolations")
+ .option("snapshotPartitionId", partition_id)
+ .option("snapshotStartBatchId", 0)
+ .load()
+ )
+
+ assert (
+ value_state_df.select("partition_id", "groupingKey", "value")
+ .filter(value_state_df["partition_id"] == partition_id)
+ .sort("groupingKey")
+ .collect()
+ == state_snapshot_df.selectExpr(
+ "partition_id", "key.id AS groupingKey", "value.value AS value"
+ )
+ .sort("groupingKey")
+ .collect()
+ )
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ with self.sql_conf(
+ {"spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled": "true"}
+ ):
+ self._test_transform_with_state_in_pandas_basic(
+ SimpleStatefulProcessor(),
+ check_results,
+ False,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ )
+
+ """
Review Comment:
Sorry, typo.. Already removed.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1901991247
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +223,108 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+/**
+ * TransformWithStateInPandas driver side Python runner. Similar as executor side runner,
+ * will start a new daemon thread on the Python runner to run state server.
+ */
+class TransformWithStateInPandasPythonPreInitRunner(
+ func: PythonFunction,
+ workerModule: String,
+ timeZoneId: String,
+ groupingKeySchema: StructType,
+ processorHandleImpl: DriverStatefulProcessorHandleImpl)
+ extends StreamingPythonRunner(func, "", "", workerModule)
+ with TransformWithStateInPandasPythonRunnerUtils
+ with Logging {
+ protected val sqlConf = SQLConf.get
+
+ private var dataOut: DataOutputStream = _
+ private var dataIn: DataInputStream = _
+
+ private var daemonThread: Thread = _
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ val result = super.init()
+ dataOut = result._1
+ dataIn = result._2
+
+ // start state server, update socket port
+ startStateServer()
+ (dataOut, dataIn)
+ }
+
+ def process(): Unit = {
+ // Also write the port number for state server
+ dataOut.writeInt(stateServerSocketPort)
+ PythonWorkerUtils.writeUTF(groupingKeySchema.json, dataOut)
+ dataOut.flush()
+
+ val resFromPython = dataIn.readInt()
+ if (resFromPython != 0) {
+ val errMessage = PythonWorkerUtils.readUTF(dataIn)
+ throw streamingPythonRunnerInitializationFailure(resFromPython, errMessage)
+ }
+ }
+
+ override def stop(): Unit = {
+ super.stop()
+ closeServerSocketChannelSilently(stateServerSocket)
+ daemonThread.stop()
Review Comment:
We shouldn't continue to use the `Thread.stop()` method, because in Java 21 it will directly throw an `UnsupportedOperationException`:
https://github.com/openjdk/jdk/blob/890adb6410dab4606a4f26a942aed02fb2f55387/src/java.base/share/classes/java/lang/Thread.java#L1665-L1668
this issue has caused the [daily test in Java 21](https://github.com/apache/spark/actions/workflows/build_java21.yml) to fail starting from 7 days ago. Can you please address this issue? @jingz-db
- https://github.com/apache/spark/actions/runs/12511573912/job/34903859772
- https://github.com/apache/spark/actions/runs/12523542188/job/34933207012
- https://github.com/apache/spark/actions/runs/12592534465/job/35097321533
also cc @HeartSaVioR and @dongjoon-hyun
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1901991247
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +223,108 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+/**
+ * TransformWithStateInPandas driver side Python runner. Similar as executor side runner,
+ * will start a new daemon thread on the Python runner to run state server.
+ */
+class TransformWithStateInPandasPythonPreInitRunner(
+ func: PythonFunction,
+ workerModule: String,
+ timeZoneId: String,
+ groupingKeySchema: StructType,
+ processorHandleImpl: DriverStatefulProcessorHandleImpl)
+ extends StreamingPythonRunner(func, "", "", workerModule)
+ with TransformWithStateInPandasPythonRunnerUtils
+ with Logging {
+ protected val sqlConf = SQLConf.get
+
+ private var dataOut: DataOutputStream = _
+ private var dataIn: DataInputStream = _
+
+ private var daemonThread: Thread = _
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ val result = super.init()
+ dataOut = result._1
+ dataIn = result._2
+
+ // start state server, update socket port
+ startStateServer()
+ (dataOut, dataIn)
+ }
+
+ def process(): Unit = {
+ // Also write the port number for state server
+ dataOut.writeInt(stateServerSocketPort)
+ PythonWorkerUtils.writeUTF(groupingKeySchema.json, dataOut)
+ dataOut.flush()
+
+ val resFromPython = dataIn.readInt()
+ if (resFromPython != 0) {
+ val errMessage = PythonWorkerUtils.readUTF(dataIn)
+ throw streamingPythonRunnerInitializationFailure(resFromPython, errMessage)
+ }
+ }
+
+ override def stop(): Unit = {
+ super.stop()
+ closeServerSocketChannelSilently(stateServerSocket)
+ daemonThread.stop()
Review Comment:
We shouldn't continue to use the `Thread.stop()` method, because in Java 21 it will directly throw an `UnsupportedOperationException`:
https://github.com/openjdk/jdk/blob/890adb6410dab4606a4f26a942aed02fb2f55387/src/java.base/share/classes/java/lang/Thread.java#L1665-L1668
this issue has caused the daily test in Java 21 to fail starting from 7 days ago. Can you please address this issue? @jingz-db
- https://github.com/apache/spark/actions/runs/12511573912/job/34903859772
- https://github.com/apache/spark/actions/runs/12523542188/job/34933207012
- https://github.com/apache/spark/actions/runs/12592534465/job/35097321533
cc @HeartSaVioR
also cc @dongjoon-hyun
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1889346389
##########
python/pyspark/sql/streaming/transform_with_state_driver_worker.py:
##########
@@ -0,0 +1,89 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import json
+
+from pyspark.util import local_connect_and_auth
+from pyspark.serializers import (
+ write_int,
+ read_int,
+ UTF8Deserializer,
+ CPickleSerializer,
+)
+from pyspark import worker
+from pyspark.util import handle_worker_exception
+from typing import IO
+from pyspark.worker_util import check_python_version
+from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
+from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+from pyspark.sql.types import StructType
+
+pickle_ser = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def main(infile: IO, outfile: IO) -> None:
+ check_python_version(infile)
+
+ log_name = "Streaming TransformWithStateInPandas Python worker"
+
+ def process(processor, mode, key, input):
+ print(f"{log_name} Started execution of UDF: {func}.\n")
+ func(processor, mode, key, input)
+ print(f"{log_name} Completed execution of UDF: {func}.\n")
+
+ try:
+ func, return_type = worker.read_command(pickle_ser, infile)
+ # send signal for getting args
+ write_int(0, outfile)
+ outfile.flush()
+
+ while True:
+ state_server_port = read_int(infile)
+ key_schema = StructType.fromJson(json.loads(utf8_deserializer.loads(infile)))
+ print(
+ f"{log_name} received parameters for UDF. State server port: {state_server_port}, "
+ f"key schema: {key_schema}.\n"
+ )
+
+ stateful_processor_api_client = StatefulProcessorApiClient(
+ state_server_port, key_schema
+ )
+ process(
+ stateful_processor_api_client,
+ TransformWithStateInPandasFuncMode.PRE_INIT,
+ None,
+ iter([]),
+ )
+ write_int(0, outfile)
+ outfile.flush()
+ except Exception as e:
+ handle_worker_exception(e, outfile)
+ outfile.flush()
+
+
+if __name__ == "__main__":
+ # Read information about how to connect back to the JVM from the environment.
+ java_port = int(os.environ["PYTHON_WORKER_FACTORY_PORT"])
+ auth_secret = os.environ["PYTHON_WORKER_FACTORY_SECRET"]
+ (sock_file, sock) = local_connect_and_auth(java_port, auth_secret)
+ # There could be a long time between each micro batch.
+ sock.settimeout(None)
Review Comment:
Sure, sg. I am deleting the code of explicitly setting timeout and just use the default like the worker on executor.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1893860278
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasExec.scala:
##########
@@ -109,6 +112,21 @@ case class TransformWithStateInPandasExec(
// Each state variable has its own schema, this is a dummy one.
protected val schemaForValueRow: StructType = new StructType().add("value", BinaryType)
+ override def operatorStateMetadataVersion: Int = 2
+
+ override def getColFamilySchemas(): Map[String, StateStoreColFamilySchema] = {
+ driverProcessorHandle.getColumnFamilySchemas
+ }
+
+ override def getStateVariableInfos(): Map[String, TransformWithStateVariableInfo] = {
+ driverProcessorHandle.getStateVariableInfos
+ }
+
+ /** Metadata of this stateful operator and its states stores.
+ * Written during IncrementalExecution */
Review Comment:
nit: Let's explicitly mention which method will initialize this - looks like driverProcessorHandle would be useful "after validateAndMaybeEvolveStateSchema is called".
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +224,137 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+/**
+ * TransformWithStateInPandas driver side Python runner. Similar as executor side runner,
+ * will start a new daemon thread on the Python runner to run state server.
+ */
+class TransformWithStateInPandasPythonPreInitRunner(
Review Comment:
nit: This class has a lot of duplications - I'm OK to leave it as it is as of now, but let's consider it as tech debt.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +224,137 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+/**
+ * TransformWithStateInPandas driver side Python runner. Similar as executor side runner,
+ * will start a new daemon thread on the Python runner to run state server.
+ */
+class TransformWithStateInPandasPythonPreInitRunner(
+ func: PythonFunction,
+ workerModule: String,
+ timeZoneId: String,
+ groupingKeySchema: StructType,
+ processorHandleImpl: DriverStatefulProcessorHandleImpl)
+ extends StreamingPythonRunner(func, "", "", workerModule)
+ with TransformWithStateInPandasPythonRunnerUtils
+ with Logging {
+ protected val sqlConf = SQLConf.get
+
+ private var dataIn: DataInputStream = _
+ private var dataOut: DataOutputStream = _
+
+ private var daemonThread: Thread = _
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ val env = SparkEnv.get
+
+ val localdir = env.blockManager.diskBlockManager.localDirs.map(f => f.getPath()).mkString(",")
+ envVars.put("SPARK_LOCAL_DIRS", localdir)
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ val workerFactory =
+ new PythonWorkerFactory(pythonExec, workerModule, envVars.asScala.toMap, false)
+ val (worker: PythonWorker, _) = workerFactory.createSimpleWorker(blockingMode = true)
+ pythonWorker = Some(worker)
+ pythonWorkerFactory = Some(workerFactory)
+
+ val stream = new BufferedOutputStream(
+ pythonWorker.get.channel.socket().getOutputStream, bufferSize)
+ dataOut = new DataOutputStream(stream)
+
+ PythonWorkerUtils.writePythonVersion(pythonVer, dataOut)
+
+ // Send the user function to python process
+ PythonWorkerUtils.writePythonFunction(func, dataOut)
+ dataOut.flush()
+
+ dataIn = new DataInputStream(
+ new BufferedInputStream(pythonWorker.get.channel.socket().getInputStream, bufferSize))
+
+ val resFromPython = dataIn.readInt()
+ if (resFromPython != 0) {
+ val errMessage = PythonWorkerUtils.readUTF(dataIn)
+ throw streamingPythonRunnerInitializationFailure(resFromPython, errMessage)
+ }
+ logInfo("Driver Python Runner initialization succeeded.")
+
+ // start state server, update socket port
+ startStateServer()
+ (dataOut, dataIn)
+ }
+
+ def process(): Unit = {
+ // Also write the port number for state server
+ dataOut.writeInt(stateServerSocketPort)
+ PythonWorkerUtils.writeUTF(groupingKeySchema.json, dataOut)
+ dataOut.flush()
+
+ val resFromPython = dataIn.readInt()
+ if (resFromPython != 0) {
+ val errMessage = PythonWorkerUtils.readUTF(dataIn)
+ throw streamingPythonRunnerInitializationFailure(resFromPython, errMessage)
+ }
+ }
+
+ override def stop(): Unit = {
+ super.stop()
+ closeServerSocketChannelSilently(stateServerSocket)
+ daemonThread.stop()
+ }
+
+ private def startStateServer(): Unit = {
+ initStateServer()
+
+ daemonThread = new Thread {
+ override def run(): Unit = {
+ try {
+ new TransformWithStateInPandasStateServer(stateServerSocket, processorHandleImpl,
+ groupingKeySchema, timeZoneId, errorOnDuplicatedFieldNames = true,
+ largeVarTypes = sqlConf.arrowUseLargeVarTypes,
+ sqlConf.arrowTransformWithStateInPandasMaxRecordsPerBatch).run()
+ } catch {
+ case e: Exception =>
+ throw new SparkException("TransformWithStateInPandas state server " +
+ "daemon thread exited unexpectedly (crashed)", e)
+ }
+ }
+ }
+ daemonThread.setDaemon(true)
+ daemonThread.setName("stateConnectionListenerThread")
+ daemonThread.start()
+ }
+}
+
+/**
+ * TransformWithStateInPandas Python runner utils functions for handling a state server
+ * in a new daemon thread.
+ */
+trait TransformWithStateInPandasPythonRunnerUtils extends Logging{
Review Comment:
nit: space between `g` and `{`
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -358,7 +369,7 @@ def check_results(batch_df, batch_id):
for q in self.spark.streams.active:
q.stop()
if batch_id == 0 or batch_id == 1:
- time.sleep(6)
+ time.sleep(4)
Review Comment:
nit: Is it relevant change? I am happy with reducing test duration, just to see the rationale.
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -852,13 +872,6 @@ def check_results(batch_df, batch_id):
StatefulProcessorWithInitialStateTimers(), check_results, "processingTime"
)
- # run the same test suites again but with single shuffle partition
Review Comment:
Isn't this something we intentionally added because of some observed failure? Mind explaining why we are going to remove this?
##########
python/pyspark/sql/streaming/transform_with_state_driver_worker.py:
##########
@@ -0,0 +1,103 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import json
+from typing import Any, Iterator, TYPE_CHECKING
+
+from pyspark.util import local_connect_and_auth
+from pyspark.serializers import (
+ write_int,
+ read_int,
+ UTF8Deserializer,
+ CPickleSerializer,
+)
+from pyspark import worker
+from pyspark.util import handle_worker_exception
+from typing import IO
+from pyspark.worker_util import check_python_version
+from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
+from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+from pyspark.sql.types import StructType
+
+if TYPE_CHECKING:
+ from pyspark.sql.pandas._typing import (
+ DataFrameLike as PandasDataFrameLike,
+ )
+
+pickle_ser = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def main(infile: IO, outfile: IO) -> None:
+ check_python_version(infile)
+
+ log_name = "Streaming TransformWithStateInPandas Python worker"
+ print(f"Starting {log_name}.\n")
+
+ def process(
+ processor: StatefulProcessorApiClient,
+ mode: TransformWithStateInPandasFuncMode,
+ key: Any,
+ input: Iterator["PandasDataFrameLike"],
+ ) -> None:
+ print(f"{log_name} Starting execution of UDF: {func}.\n")
+ func(processor, mode, key, input)
+ print(f"{log_name} Completed execution of UDF: {func}.\n")
+
+ try:
+ func, return_type = worker.read_command(pickle_ser, infile)
+ print(
+ f"{log_name} finish init stage of Python runner. Received UDF from JVM: {func}, "
+ f"received return type of UDF: {return_type}.\n"
+ )
+ # send signal for getting args
+ write_int(0, outfile)
+ outfile.flush()
+
+ while True:
Review Comment:
Is this intentional? Is this worker process reused by multiple queries in the same driver? Otherwise how this would be used?
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +920,359 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ # This test covers mapState with TTL, an empty state variable
+ # and additional test against initial state python runner
+ def test_transform_with_map_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
+ map_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "mapState")
+ .load()
+ )
+ assert set(
+ map_state_df.selectExpr(
+ "key.id AS groupingKey",
+ "user_map_key.name AS mapKey",
+ "user_map_value.value.count AS mapValue",
+ )
+ .sort("groupingKey")
+ .collect()
+ ) == {
+ Row(groupingKey="0", mapKey="key2", mapValue=2),
+ Row(groupingKey="1", mapKey="key2", mapValue=2),
+ }
+
+ ttl_df = map_state_df.selectExpr(
+ "user_map_value.ttlExpirationMs AS TTLVal"
+ ).collect()
+ # check if there are two rows containing TTL value in map state dataframe
+ assert len(ttl_df) == 2
+ # check if two rows are of the same TTL value
+ assert len(set(ttl_df)) == 1
+
+ list_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState")
+ .load()
+ )
+ assert list_state_df.isEmpty()
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # run the same test suite again but with no-op initial state
+ # TWS with initial state is using a different python runner
+ init_data = [("0", 789), ("3", 987)]
+ initial_state = self.spark.createDataFrame(init_data, "id string, temperature int").groupBy(
+ "id"
+ )
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=initial_state,
+ )
+
+ # This test covers multiple list state variables and flatten option
+ def test_transform_with_list_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] in ["listState1", "listState2"]:
+ state_var["stateVariableType"] == "ListState"
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source and flatten option
+ list_state_1_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState1")
+ .option("flattenCollectionTypes", True)
+ .load()
+ )
+ assert list_state_1_df.selectExpr(
+ "key.id AS groupingKey",
+ "list_element.temperature AS listElement",
+ ).sort("groupingKey", "listElement").collect() == [
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=111),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=111),
+ Row(groupingKey="1", listElement=120),
+ Row(groupingKey="1", listElement=120),
+ ]
+
+ list_state_2_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState2")
+ .option("flattenCollectionTypes", False)
+ .load()
+ )
+ assert list_state_2_df.selectExpr(
+ "key.id AS groupingKey", "list_value.temperature AS valueList"
+ ).sort("groupingKey").withColumn(
+ "valueSortedList", array_sort(col("valueList"))
+ ).select(
+ "groupingKey", "valueSortedList"
+ ).collect() == [
+ Row(groupingKey="0", valueSortedList=[20, 20, 120, 120, 222]),
+ Row(groupingKey="1", valueSortedList=[20, 20, 120, 120, 222]),
+ ]
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ ListStateProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # This test covers value state variable and read change feed,
+ # snapshotStartBatchId related options
+ def test_transform_with_value_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="3"),
+ Row(id="1", countAsString="2"),
+ }
+
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] in ["numViolations", "tempState"]:
+ state_var["stateVariableType"] == "ValueState"
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source and readChangeFeed
+ value_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "numViolations")
+ .option("readChangeFeed", True)
+ .option("changeStartBatchId", 0)
+ .load()
+ ).selectExpr(
+ "change_type", "key.id AS groupingKey", "value.value AS value", "partition_id"
+ )
+
+ assert value_state_df.select("change_type", "groupingKey", "value").sort(
+ "groupingKey"
+ ).collect() == [
+ Row(change_type="update", groupingKey="0", value=1),
+ Row(change_type="update", groupingKey="1", value=2),
+ ]
+
+ partition_id_list = [
+ row["partition_id"] for row in value_state_df.select("partition_id").collect()
+ ]
+
+ for partition_id in partition_id_list:
+ # check for state data source and snapshotStartBatchId options
+ state_snapshot_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "numViolations")
+ .option("snapshotPartitionId", partition_id)
+ .option("snapshotStartBatchId", 0)
+ .load()
+ )
+
+ assert (
+ value_state_df.select("partition_id", "groupingKey", "value")
+ .filter(value_state_df["partition_id"] == partition_id)
+ .sort("groupingKey")
+ .collect()
+ == state_snapshot_df.selectExpr(
+ "partition_id", "key.id AS groupingKey", "value.value AS value"
+ )
+ .sort("groupingKey")
+ .collect()
+ )
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ with self.sql_conf(
+ {"spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled": "true"}
+ ):
+ self._test_transform_with_state_in_pandas_basic(
+ SimpleStatefulProcessor(),
+ check_results,
+ False,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ )
+
+ """
Review Comment:
nit: Is this intentional? Tests are commented out.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887728579
##########
python/pyspark/sql/streaming/transform_with_state_driver_worker.py:
##########
@@ -0,0 +1,89 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import json
+
+from pyspark.util import local_connect_and_auth
+from pyspark.serializers import (
+ write_int,
+ read_int,
+ UTF8Deserializer,
+ CPickleSerializer,
+)
+from pyspark import worker
+from pyspark.util import handle_worker_exception
+from typing import IO
+from pyspark.worker_util import check_python_version
+from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
+from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+from pyspark.sql.types import StructType
+
+pickle_ser = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def main(infile: IO, outfile: IO) -> None:
+ check_python_version(infile)
+
+ log_name = "Streaming TransformWithStateInPandas Python worker"
+
+ def process(processor, mode, key, input):
+ print(f"{log_name} Started execution of UDF: {func}.\n")
+ func(processor, mode, key, input)
+ print(f"{log_name} Completed execution of UDF: {func}.\n")
+
+ try:
+ func, return_type = worker.read_command(pickle_ser, infile)
+ # send signal for getting args
+ write_int(0, outfile)
+ outfile.flush()
+
+ while True:
+ state_server_port = read_int(infile)
+ key_schema = StructType.fromJson(json.loads(utf8_deserializer.loads(infile)))
+ print(
+ f"{log_name} received parameters for UDF. State server port: {state_server_port}, "
+ f"key schema: {key_schema}.\n"
+ )
+
+ stateful_processor_api_client = StatefulProcessorApiClient(
+ state_server_port, key_schema
+ )
+ process(
+ stateful_processor_api_client,
+ TransformWithStateInPandasFuncMode.PRE_INIT,
+ None,
+ iter([]),
+ )
+ write_int(0, outfile)
+ outfile.flush()
+ except Exception as e:
+ handle_worker_exception(e, outfile)
+ outfile.flush()
+
+
+if __name__ == "__main__":
+ # Read information about how to connect back to the JVM from the environment.
+ java_port = int(os.environ["PYTHON_WORKER_FACTORY_PORT"])
+ auth_secret = os.environ["PYTHON_WORKER_FACTORY_SECRET"]
+ (sock_file, sock) = local_connect_and_auth(java_port, auth_secret)
+ # There could be a long time between each micro batch.
+ sock.settimeout(None)
Review Comment:
I am a bit hesitate what is a good number to set here. The timeout = overhead set up and tear down time + user defined init() function execution time + network delay. This is quite indefinite as we can not predict how long it will take to finish `init()` UDF. What do you think? Do you think we should set this a config, or shall we just keep as indefinite?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887542845
##########
python/pyspark/sql/streaming/transform_with_state_driver_worker.py:
##########
@@ -0,0 +1,89 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import json
+
+from pyspark.util import local_connect_and_auth
+from pyspark.serializers import (
+ write_int,
+ read_int,
+ UTF8Deserializer,
+ CPickleSerializer,
+)
+from pyspark import worker
+from pyspark.util import handle_worker_exception
+from typing import IO
+from pyspark.worker_util import check_python_version
+from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
+from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+from pyspark.sql.types import StructType
+
+pickle_ser = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def main(infile: IO, outfile: IO) -> None:
Review Comment:
Could we add some file level comment here ?
##########
python/pyspark/sql/streaming/proto/StateMessage_pb2.py:
##########
@@ -51,7 +51,7 @@
if not _descriptor._USE_C_DESCRIPTORS:
DESCRIPTOR._loaded_options = None
_globals["_HANDLESTATE"]._serialized_start = 6408
- _globals["_HANDLESTATE"]._serialized_end = 6504
+ _globals["_HANDLESTATE"]._serialized_end = 6518
Review Comment:
this is auto-generated ?
##########
python/pyspark/sql/streaming/stateful_processor_util.py:
##########
@@ -25,3 +25,4 @@ class TransformWithStateInPandasFuncMode(Enum):
PROCESS_DATA = 1
PROCESS_TIMER = 2
COMPLETE = 3
+ PRE_INIT = 4
Review Comment:
What does this mean ?
##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -504,6 +504,20 @@ def transformWithStateInPandas(
if isinstance(outputStructType, str):
outputStructType = cast(StructType, _parse_datatype_string(outputStructType))
+ def handle_pre_init(
+ statefulProcessorApiClient: StatefulProcessorApiClient,
+ ) -> Iterator["PandasDataFrameLike"]:
+ # driver handle is different from the handle used on executors
+ driver_handle = StatefulProcessorHandle(statefulProcessorApiClient)
+ statefulProcessorApiClient.set_handle_state(StatefulProcessorHandleState.PRE_INIT)
+ statefulProcessor.init(driver_handle)
+
+ # TODO figure out why we need to call close here in Scala side
Review Comment:
mostly just to reset the internal state right ?
##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -504,6 +504,20 @@ def transformWithStateInPandas(
if isinstance(outputStructType, str):
outputStructType = cast(StructType, _parse_datatype_string(outputStructType))
+ def handle_pre_init(
+ statefulProcessorApiClient: StatefulProcessorApiClient,
+ ) -> Iterator["PandasDataFrameLike"]:
+ # driver handle is different from the handle used on executors
Review Comment:
lets also briefly say why its different ?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887720909
##########
python/pyspark/sql/streaming/stateful_processor_util.py:
##########
@@ -25,3 +25,4 @@ class TransformWithStateInPandasFuncMode(Enum):
PROCESS_DATA = 1
PROCESS_TIMER = 2
COMPLETE = 3
+ PRE_INIT = 4
Review Comment:
This is a stateful processor handle state introduced only for internal use. It will be passed into our UDF `transformWithStateUDF` in `group_ops.py` as an internal parameter and we will use this parameter to decide whether we are going to process data, timer, or if we are at pre-init stage. The public exposed StatefulProcessorHandleState will only be set on the JVM.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1896091475
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +224,137 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+/**
+ * TransformWithStateInPandas driver side Python runner. Similar as executor side runner,
+ * will start a new daemon thread on the Python runner to run state server.
+ */
+class TransformWithStateInPandasPythonPreInitRunner(
Review Comment:
Made some small modifications in the base class and refactor this class to remove duplications in `init()`.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #49156:
URL: https://github.com/apache/spark/pull/49156#issuecomment-2564267751
@jingz-db
https://github.com/apache/spark/pull/47805#issuecomment-2564246240
Looks like Thread.stop() is removed in Java 21 - it was deprecated. Looks like they no longer support the same operation but recommend to use a variable as signal flag. Could you please create a follow-up PR? Even better if you could run the failing test with JDK 21 setup. Thanks in advance!
cc. @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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #49156:
URL: https://github.com/apache/spark/pull/49156#issuecomment-2563172329
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1901991247
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +223,108 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+/**
+ * TransformWithStateInPandas driver side Python runner. Similar as executor side runner,
+ * will start a new daemon thread on the Python runner to run state server.
+ */
+class TransformWithStateInPandasPythonPreInitRunner(
+ func: PythonFunction,
+ workerModule: String,
+ timeZoneId: String,
+ groupingKeySchema: StructType,
+ processorHandleImpl: DriverStatefulProcessorHandleImpl)
+ extends StreamingPythonRunner(func, "", "", workerModule)
+ with TransformWithStateInPandasPythonRunnerUtils
+ with Logging {
+ protected val sqlConf = SQLConf.get
+
+ private var dataOut: DataOutputStream = _
+ private var dataIn: DataInputStream = _
+
+ private var daemonThread: Thread = _
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ val result = super.init()
+ dataOut = result._1
+ dataIn = result._2
+
+ // start state server, update socket port
+ startStateServer()
+ (dataOut, dataIn)
+ }
+
+ def process(): Unit = {
+ // Also write the port number for state server
+ dataOut.writeInt(stateServerSocketPort)
+ PythonWorkerUtils.writeUTF(groupingKeySchema.json, dataOut)
+ dataOut.flush()
+
+ val resFromPython = dataIn.readInt()
+ if (resFromPython != 0) {
+ val errMessage = PythonWorkerUtils.readUTF(dataIn)
+ throw streamingPythonRunnerInitializationFailure(resFromPython, errMessage)
+ }
+ }
+
+ override def stop(): Unit = {
+ super.stop()
+ closeServerSocketChannelSilently(stateServerSocket)
+ daemonThread.stop()
Review Comment:
We shouldn't continue to use the `Thread.stop()` method, because in Java 21 it will directly throw an `UnsupportedOperationException`:
https://github.com/openjdk/jdk/blob/890adb6410dab4606a4f26a942aed02fb2f55387/src/java.base/share/classes/java/lang/Thread.java#L1665-L1668
this issue has caused the daily test in Java 21 to fail starting from 7 days ago. Can you please address this issue? @jingz-db
- https://github.com/apache/spark/actions/runs/12511573912/job/34903859772
- https://github.com/apache/spark/actions/runs/12523542188/job/34933207012
- https://github.com/apache/spark/actions/runs/12592534465/job/35097321533
also cc @HeartSaVioR and @dongjoon-hyun
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1889357588
##########
python/pyspark/sql/tests/pandas/test_pandas_transform_with_state.py:
##########
@@ -907,6 +923,328 @@ def test_transform_with_state_in_pandas_batch_query_initial_state(self):
Row(id="1", value=str(146 + 346)),
}
+ # This test covers mapState with TTL, an empty state variable
+ # and additional test against initial state python runner
+ def test_transform_with_map_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ assert set(
+ metadata_df.select(
+ "operatorId",
+ "operatorName",
+ "stateStoreName",
+ "numPartitions",
+ "minBatchId",
+ "maxBatchId",
+ ).collect()
+ ) == {
+ Row(
+ operatorId=0,
+ operatorName="transformWithStateInPandasExec",
+ stateStoreName="default",
+ numPartitions=5,
+ minBatchId=0,
+ maxBatchId=0,
+ )
+ }
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ assert operator_properties_json_obj["timeMode"] == "ProcessingTime"
+ assert operator_properties_json_obj["outputMode"] == "Update"
+
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] == "mapState":
+ assert state_var["stateVariableType"] == "MapState"
+ assert state_var["ttlEnabled"]
+ elif state_var["stateName"] == "listState":
+ assert state_var["stateVariableType"] == "ListState"
+ assert not state_var["ttlEnabled"]
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source
+ map_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "mapState")
+ .load()
+ )
+ assert set(
+ map_state_df.selectExpr(
+ "key.id AS groupingKey",
+ "user_map_key.name AS mapKey",
+ "user_map_value.value.count AS mapValue",
+ )
+ .sort("groupingKey")
+ .collect()
+ ) == {
+ Row(groupingKey="0", mapKey="key2", mapValue=2),
+ Row(groupingKey="1", mapKey="key2", mapValue=2),
+ }
+
+ ttl_df = map_state_df.selectExpr(
+ "user_map_value.ttlExpirationMs AS TTLVal"
+ ).collect()
+ # check if there are two rows containing TTL value in map state dataframe
+ assert len(ttl_df) == 2
+ # check if two rows are of the same TTL value
+ assert len(set(ttl_df)) == 1
+
+ list_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState")
+ .load()
+ )
+ assert list_state_df.isEmpty()
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # run the same test suite again but with no-op initial state
+ # TWS with initial state is using a different python runner
+ init_data = [("0", 789), ("3", 987)]
+ initial_state = self.spark.createDataFrame(init_data, "id string, temperature int").groupBy(
+ "id"
+ )
+ self._test_transform_with_state_in_pandas_basic(
+ MapStateLargeTTLProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=initial_state,
+ )
+
+ # This test covers multiple list state variables and flatten option
+ def test_transform_with_list_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] in ["listState1", "listState2"]:
+ state_var["stateVariableType"] == "ListState"
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source and flatten option
+ list_state_1_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState1")
+ .option("flattenCollectionTypes", True)
+ .load()
+ )
+ assert list_state_1_df.selectExpr(
+ "key.id AS groupingKey",
+ "list_element.temperature AS listElement",
+ ).sort("groupingKey", "listElement").collect() == [
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=20),
+ Row(groupingKey="0", listElement=111),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="0", listElement=120),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=20),
+ Row(groupingKey="1", listElement=111),
+ Row(groupingKey="1", listElement=120),
+ Row(groupingKey="1", listElement=120),
+ ]
+
+ list_state_2_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "listState2")
+ .option("flattenCollectionTypes", False)
+ .load()
+ )
+ assert list_state_2_df.selectExpr(
+ "key.id AS groupingKey", "list_value.temperature AS valueList"
+ ).sort("groupingKey").withColumn(
+ "valueSortedList", array_sort(col("valueList"))
+ ).select(
+ "groupingKey", "valueSortedList"
+ ).collect() == [
+ Row(groupingKey="0", valueSortedList=[20, 20, 120, 120, 222]),
+ Row(groupingKey="1", valueSortedList=[20, 20, 120, 120, 222]),
+ ]
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ self._test_transform_with_state_in_pandas_basic(
+ ListStateProcessor(),
+ check_results,
+ True,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ initial_state=None,
+ )
+
+ # This test covers value state variable and read change feed
+ def test_transform_with_value_state_metadata(self):
+ checkpoint_path = tempfile.mktemp()
+
+ def check_results(batch_df, batch_id):
+ if batch_id == 0:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="2"),
+ Row(id="1", countAsString="2"),
+ }
+ else:
+ assert set(batch_df.sort("id").collect()) == {
+ Row(id="0", countAsString="3"),
+ Row(id="1", countAsString="2"),
+ }
+
+ # check for state metadata source
+ metadata_df = self.spark.read.format("state-metadata").load(checkpoint_path)
+ operator_properties_json_obj = json.loads(
+ metadata_df.select("operatorProperties").collect()[0][0]
+ )
+ state_var_list = operator_properties_json_obj["stateVariables"]
+
+ # TODO "tempState" should not appear in the metadata as it is already deleted
+ assert len(state_var_list) == 3
+ for state_var in state_var_list:
+ if state_var["stateName"] in ["numViolations", "tempState"]:
+ state_var["stateVariableType"] == "ValueState"
+ else:
+ assert state_var["stateName"] == "$procTimers_keyToTimestamp"
+ assert state_var["stateVariableType"] == "TimerState"
+
+ # check for state data source and flatten option
+ value_state_df = (
+ self.spark.read.format("statestore")
+ .option("path", checkpoint_path)
+ .option("stateVarName", "numViolations")
+ .option("readChangeFeed", True)
+ .option("changeStartBatchId", 0)
+ .load()
+ )
+
+ assert value_state_df.selectExpr(
+ "change_type", "key.id AS groupingKey", "value.value AS value"
+ ).sort("groupingKey").collect() == [
+ Row(change_type="update", groupingKey="0", value=1),
+ Row(change_type="update", groupingKey="1", value=2),
+ ]
+
+ for q in self.spark.streams.active:
+ q.stop()
+
+ with self.sql_conf(
+ {"spark.sql.streaming.stateStore.rocksdb.changelogCheckpointing.enabled": "true"}
+ ):
+ self._test_transform_with_state_in_pandas_basic(
+ SimpleStatefulProcessor(),
+ check_results,
+ False,
+ "processingTime",
+ checkpoint_path=checkpoint_path,
+ )
+
+ def test_transform_with_state_restart_with_multiple_rows_init_state(self):
Review Comment:
Could we also add 1 test for the `startSnapshotBatchId` related options ?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887719504
##########
python/pyspark/sql/streaming/proto/StateMessage_pb2.py:
##########
@@ -51,7 +51,7 @@
if not _descriptor._USE_C_DESCRIPTORS:
DESCRIPTOR._loaded_options = None
_globals["_HANDLESTATE"]._serialized_start = 6408
- _globals["_HANDLESTATE"]._serialized_end = 6504
+ _globals["_HANDLESTATE"]._serialized_end = 6518
Review Comment:
Yes, it is auto generated from `StateMessage.proto`
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887724687
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +224,130 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+class TransformWithStateInPandasPythonPreInitRunner(
+ func: PythonFunction,
+ workerModule: String,
+ timeZoneId: String,
+ groupingKeySchema: StructType,
+ processorHandleImpl: DriverStatefulProcessorHandleImpl)
+ extends StreamingPythonRunner(func, "", "", workerModule)
+ with TransformWithStateInPandasPythonRunnerUtils
+ with Logging {
+ protected val sqlConf = SQLConf.get
+
+ private var dataIn: DataInputStream = _
+ private var dataOut: DataOutputStream = _
+
+ private var daemonThread: Thread = _
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ val env = SparkEnv.get
+
+ val localdir = env.blockManager.diskBlockManager.localDirs.map(f => f.getPath()).mkString(",")
Review Comment:
I am not one hundred percent sure but this is exactly same as what the base `StreamingPythonRunner` class is doing here: https://github.com/apache/spark/blob/master/core/src/main/scala/org/apache/spark/api/python/StreamingPythonRunner.scala#L66 and also the same as custom python source here: https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonStreamingSourceRunner.scala#L89
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887719936
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +224,130 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+class TransformWithStateInPandasPythonPreInitRunner(
Review Comment:
lets add a class level comment ?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887720557
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/IncrementalExecution.scala:
##########
@@ -217,6 +217,9 @@ class IncrementalExecution(
case _: TransformWithStateExec =>
sparkSession.sessionState.conf.
getConf(SQLConf.STREAMING_TRANSFORM_WITH_STATE_OP_STATE_SCHEMA_VERSION)
+ case _: TransformWithStateInPandasExec =>
Review Comment:
Could we combine into single 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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1887720906
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/TransformWithStateVariableUtils.scala:
##########
@@ -158,3 +163,98 @@ object TransformWithStateOperatorProperties extends Logging {
}
}
}
+
+trait TransformWithStateMetadataUtils extends Logging {
Review Comment:
lets add a comment here too ?
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "jingz-db (via GitHub)" <gi...@apache.org>.
jingz-db commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1896071845
##########
python/pyspark/sql/streaming/transform_with_state_driver_worker.py:
##########
@@ -0,0 +1,103 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import json
+from typing import Any, Iterator, TYPE_CHECKING
+
+from pyspark.util import local_connect_and_auth
+from pyspark.serializers import (
+ write_int,
+ read_int,
+ UTF8Deserializer,
+ CPickleSerializer,
+)
+from pyspark import worker
+from pyspark.util import handle_worker_exception
+from typing import IO
+from pyspark.worker_util import check_python_version
+from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
+from pyspark.sql.streaming.stateful_processor_util import TransformWithStateInPandasFuncMode
+from pyspark.sql.types import StructType
+
+if TYPE_CHECKING:
+ from pyspark.sql.pandas._typing import (
+ DataFrameLike as PandasDataFrameLike,
+ )
+
+pickle_ser = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def main(infile: IO, outfile: IO) -> None:
+ check_python_version(infile)
+
+ log_name = "Streaming TransformWithStateInPandas Python worker"
+ print(f"Starting {log_name}.\n")
+
+ def process(
+ processor: StatefulProcessorApiClient,
+ mode: TransformWithStateInPandasFuncMode,
+ key: Any,
+ input: Iterator["PandasDataFrameLike"],
+ ) -> None:
+ print(f"{log_name} Starting execution of UDF: {func}.\n")
+ func(processor, mode, key, input)
+ print(f"{log_name} Completed execution of UDF: {func}.\n")
+
+ try:
+ func, return_type = worker.read_command(pickle_ser, infile)
+ print(
+ f"{log_name} finish init stage of Python runner. Received UDF from JVM: {func}, "
+ f"received return type of UDF: {return_type}.\n"
+ )
+ # send signal for getting args
+ write_int(0, outfile)
+ outfile.flush()
+
+ while True:
Review Comment:
You are correct! Already removed `while` and leave a comment saying this driver will only be used once for each query run. This driver runner should only be used once for each query run in the first batch and should not be reused.
--
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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1902271107
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +223,108 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+/**
+ * TransformWithStateInPandas driver side Python runner. Similar as executor side runner,
+ * will start a new daemon thread on the Python runner to run state server.
+ */
+class TransformWithStateInPandasPythonPreInitRunner(
+ func: PythonFunction,
+ workerModule: String,
+ timeZoneId: String,
+ groupingKeySchema: StructType,
+ processorHandleImpl: DriverStatefulProcessorHandleImpl)
+ extends StreamingPythonRunner(func, "", "", workerModule)
+ with TransformWithStateInPandasPythonRunnerUtils
+ with Logging {
+ protected val sqlConf = SQLConf.get
+
+ private var dataOut: DataOutputStream = _
+ private var dataIn: DataInputStream = _
+
+ private var daemonThread: Thread = _
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ val result = super.init()
+ dataOut = result._1
+ dataIn = result._2
+
+ // start state server, update socket port
+ startStateServer()
+ (dataOut, dataIn)
+ }
+
+ def process(): Unit = {
+ // Also write the port number for state server
+ dataOut.writeInt(stateServerSocketPort)
+ PythonWorkerUtils.writeUTF(groupingKeySchema.json, dataOut)
+ dataOut.flush()
+
+ val resFromPython = dataIn.readInt()
+ if (resFromPython != 0) {
+ val errMessage = PythonWorkerUtils.readUTF(dataIn)
+ throw streamingPythonRunnerInitializationFailure(resFromPython, errMessage)
+ }
+ }
+
+ override def stop(): Unit = {
+ super.stop()
+ closeServerSocketChannelSilently(stateServerSocket)
+ daemonThread.stop()
Review Comment:
Thank you @jingz-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-50578][PYTHON][SS] Add support for new version of state metadata for TransformWithStateInPandas [spark]
Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #49156:
URL: https://github.com/apache/spark/pull/49156#discussion_r1902970871
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasPythonRunner.scala:
##########
@@ -238,7 +223,108 @@ abstract class TransformWithStateInPandasPythonBaseRunner[I](
super.compute(inputIterator, partitionIndex, context)
}
- private def closeServerSocketChannelSilently(stateServerSocket: ServerSocket): Unit = {
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets, None)
+ }
+}
+
+/**
+ * TransformWithStateInPandas driver side Python runner. Similar as executor side runner,
+ * will start a new daemon thread on the Python runner to run state server.
+ */
+class TransformWithStateInPandasPythonPreInitRunner(
+ func: PythonFunction,
+ workerModule: String,
+ timeZoneId: String,
+ groupingKeySchema: StructType,
+ processorHandleImpl: DriverStatefulProcessorHandleImpl)
+ extends StreamingPythonRunner(func, "", "", workerModule)
+ with TransformWithStateInPandasPythonRunnerUtils
+ with Logging {
+ protected val sqlConf = SQLConf.get
+
+ private var dataOut: DataOutputStream = _
+ private var dataIn: DataInputStream = _
+
+ private var daemonThread: Thread = _
+
+ override def init(): (DataOutputStream, DataInputStream) = {
+ val result = super.init()
+ dataOut = result._1
+ dataIn = result._2
+
+ // start state server, update socket port
+ startStateServer()
+ (dataOut, dataIn)
+ }
+
+ def process(): Unit = {
+ // Also write the port number for state server
+ dataOut.writeInt(stateServerSocketPort)
+ PythonWorkerUtils.writeUTF(groupingKeySchema.json, dataOut)
+ dataOut.flush()
+
+ val resFromPython = dataIn.readInt()
+ if (resFromPython != 0) {
+ val errMessage = PythonWorkerUtils.readUTF(dataIn)
+ throw streamingPythonRunnerInitializationFailure(resFromPython, errMessage)
+ }
+ }
+
+ override def stop(): Unit = {
+ super.stop()
+ closeServerSocketChannelSilently(stateServerSocket)
+ daemonThread.stop()
Review Comment:
I provided a simple fix to restore the daily tests for Java 21 first: https://github.com/apache/spark/pull/49354
If @jingz-db has a better solution, please feel free to submit new 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