You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "zifeif2 (via GitHub)" <gi...@apache.org> on 2026/01/28 00:37:58 UTC
[PR] [SPARK-55246][SS] Add Test for Pyspark TWS and TWSInPandas and Fix StatePartitionAllColumnFamiliesWriter Bug [spark]
zifeif2 opened a new pull request, #54019:
URL: https://github.com/apache/spark/pull/54019
<!--
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.
-->
1. Added Python tests for streaming repartition with transformWithState and transformWithStateInPandas:
2. Fixed writer bug in StatePartitionWriter.scala:
- Added checkIfUseColumnFamily() method to properly detect when column families are in use
- Removed special handling for JoinV3 operator and related assertions
- Simplified constructor by removing operatorName and sqlConf parameters
- Removed requirement for default column family to always be present
3. Fix reader bug in StatePartitionReader.scala:
- Added the same checkIfUseColumnFamily() method for consistency with writer
4. Simplified StateRewriter.scala:
- Removed SQLConf serialization/deserialization logic (no longer needed)
- Removed operatorName and sqlConf from writer instantiation
5. Removed obsolete test from StatePartitionAllColumnFamiliesWriterSuite.scala:
- Deleted test that was checking for default column family requirement (no longer enforced)
### 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.
-->
In Python TWS operators, it won't contains DEFAULT column family like Scala by default. Therefore, the original logic of detecting if an operator is using column family by checking schema size > 1 will fail on python TWS operator that only contains 1 state variable. Instead, we say an operator is using column family as long as it has a non-DEFAULT column family
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as new features, bug fixes, or other behavior changes. Documentation-only updates are not considered user-facing changes.
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'.
-->
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.
-->
See added python test.
Before fixing the writer, it failed with error:
```
pyspark.errors.exceptions.captured.UnsupportedOperationException: [STATE_STORE_UNSUPPORTED_OPERATION] iterator operation not supported with multiple column families is disabled in RocksDBStateStoreProvider SQLSTATE: XXKST
```
### 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.
-->
Yes
--
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-55246][SS] Add Test for Pyspark TWS and TWSInPandas and Fix StatePartitionAllColumnFamiliesWriter Bug [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #54019:
URL: https://github.com/apache/spark/pull/54019#discussion_r2734376997
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StatePartitionWriter.scala:
##########
@@ -55,24 +53,16 @@ class StatePartitionAllColumnFamiliesWriter(
storeName: String,
currentBatchId: Long,
colFamilyToWriterInfoMap: Map[String, StatePartitionWriterColumnFamilyInfo],
- operatorName: String,
- schemaProviderOpt: Option[StateSchemaProvider],
- sqlConf: SQLConf) {
-
- private def isJoinV3Operator(
- operatorName: String, sqlConf: SQLConf): Boolean = {
- operatorName == StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME &&
- sqlConf.getConf(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION) == 3
- }
+ schemaProviderOpt: Option[StateSchemaProvider]) {
+ // Using the heuristic that all operators that enable column families
+ // has a non-default column family
+ private val useColumnFamilies = colFamilyToWriterInfoMap.keys.toSeq
+ .exists(_ != StateStoreId.DEFAULT_STORE_NAME)
private val defaultSchema = {
colFamilyToWriterInfoMap.get(StateStore.DEFAULT_COL_FAMILY_NAME) match {
case Some(info) => info.schema
case None =>
- // joinV3 operator doesn't have default column family schema
- assert(isJoinV3Operator(operatorName, sqlConf),
Review Comment:
We are not doing `assert(useColumnFamilies)` here because `useColumnFamilies` is always true as long as there is a non-default column family.
--
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-55246][SS] Add Test for Pyspark TWS and TWSInPandas and Fix StatePartitionAllColumnFamiliesWriter Bug [spark]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #54019:
URL: https://github.com/apache/spark/pull/54019#issuecomment-3808307907
## JIRA Issue Information
=== Bug SPARK-55246 ===
Summary: Test repartition for Python TWS and TWSInPandas
Assignee: None
Status: Open
Affected: ["4.1.1"]
---
*This comment was automatically generated by GitHub Actions*
--
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-55246][SS] Add Test for Pyspark TWS and TWSInPandas and Fix StatePartitionAllColumnFamiliesWriter Bug [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on code in PR #54019:
URL: https://github.com/apache/spark/pull/54019#discussion_r2739276299
##########
python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py:
##########
@@ -320,6 +326,212 @@ def verify_after_decrease(results):
verify_after_decrease=verify_after_decrease,
)
+ def _run_tws_repartition_test(self, is_pandas):
+ """Helper method to run repartition test with a given processor factory method"""
+
+ def create_streaming_df(df):
+ # Parse text input format "id,temperature" into structured columns
+ split_df = split(df["value"], ",")
+ parsed_df = df.select(
+ split_df.getItem(0).alias("id"),
+ split_df.getItem(1).cast("integer").alias("temperature"),
+ )
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value", StringType(), True),
+ ]
+ )
+ processor = (
+ SimpleStatefulProcessorWithInitialStateFactory().pandas()
+ if is_pandas
+ else SimpleStatefulProcessorWithInitialStateFactory().row()
+ )
+ return (
+ parsed_df.groupBy("id").transformWithStateInPandas(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ if is_pandas
+ else parsed_df.groupBy("id").transformWithState(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ )
+
+ def verify_initial(results):
+ # SimpleStatefulProcessorWithInitialState accumulates temperature values
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "270", "a: 120 + 150 = 270")
+ self.assertEqual(values.get("b"), "50", "b: 50")
+ self.assertEqual(values.get("c"), "30", "c: 30")
+
+ def verify_after_increase(results):
+ # After repartition, state should be preserved and accumulated
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "371", "State for 'a': 270 + 101 = 371")
+ self.assertEqual(values.get("b"), "152", "State for 'b': 50 + 102 = 152")
+ self.assertEqual(values.get("d"), "103", "New key 'd' should have value 103")
+
+ def verify_after_decrease(results):
+ # After repartition, state should still be preserved
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "475", "State for 'a': 371 + 104 = 475")
+ self.assertEqual(values.get("c"), "135", "State for 'c': 30 + 105 = 135")
+ self.assertEqual(values.get("e"), "106", "New key 'e' should have value 106")
+
+ OfflineStateRepartitionTestUtils.run_repartition_test(
+ spark=self.spark,
+ num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
+ create_streaming_df=create_streaming_df,
+ output_mode="update",
+ batch1_data="a,120\na,150\nb,50\nc,30\n", # a:270, b:50, c:30
+ batch2_data="a,101\nb,102\nd,103\n", # a:371, b:152, d:103 (new)
+ batch3_data="a,104\nc,105\ne,106\n", # a:475, c:135, e:106 (new)
+ verify_initial=verify_initial,
+ verify_after_increase=verify_after_increase,
+ verify_after_decrease=verify_after_decrease,
+ )
+
+ def test_repartition_with_streaming_tws(self):
+ """Test repartition for streaming transformWithState"""
+
+ self._run_tws_repartition_test(is_pandas=False)
+
+ @unittest.skipIf(
+ not have_pandas or not have_pyarrow,
Review Comment:
you didn't add everything in the link i sent previously. You missed GIL: https://github.com/apache/spark/blob/9ce067c7d532a6766d622e55164baa942a065250/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_state_variable.py#L979
##########
python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py:
##########
@@ -320,6 +326,212 @@ def verify_after_decrease(results):
verify_after_decrease=verify_after_decrease,
)
+ def _run_tws_repartition_test(self, is_pandas):
+ """Helper method to run repartition test with a given processor factory method"""
+
+ def create_streaming_df(df):
+ # Parse text input format "id,temperature" into structured columns
+ split_df = split(df["value"], ",")
+ parsed_df = df.select(
+ split_df.getItem(0).alias("id"),
+ split_df.getItem(1).cast("integer").alias("temperature"),
+ )
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value", StringType(), True),
+ ]
+ )
+ processor = (
+ SimpleStatefulProcessorWithInitialStateFactory().pandas()
+ if is_pandas
+ else SimpleStatefulProcessorWithInitialStateFactory().row()
+ )
+ return (
+ parsed_df.groupBy("id").transformWithStateInPandas(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ if is_pandas
+ else parsed_df.groupBy("id").transformWithState(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ )
+
+ def verify_initial(results):
+ # SimpleStatefulProcessorWithInitialState accumulates temperature values
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "270", "a: 120 + 150 = 270")
+ self.assertEqual(values.get("b"), "50", "b: 50")
+ self.assertEqual(values.get("c"), "30", "c: 30")
+
+ def verify_after_increase(results):
+ # After repartition, state should be preserved and accumulated
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "371", "State for 'a': 270 + 101 = 371")
+ self.assertEqual(values.get("b"), "152", "State for 'b': 50 + 102 = 152")
+ self.assertEqual(values.get("d"), "103", "New key 'd' should have value 103")
+
+ def verify_after_decrease(results):
+ # After repartition, state should still be preserved
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "475", "State for 'a': 371 + 104 = 475")
+ self.assertEqual(values.get("c"), "135", "State for 'c': 30 + 105 = 135")
+ self.assertEqual(values.get("e"), "106", "New key 'e' should have value 106")
+
+ OfflineStateRepartitionTestUtils.run_repartition_test(
+ spark=self.spark,
+ num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
+ create_streaming_df=create_streaming_df,
+ output_mode="update",
+ batch1_data="a,120\na,150\nb,50\nc,30\n", # a:270, b:50, c:30
+ batch2_data="a,101\nb,102\nd,103\n", # a:371, b:152, d:103 (new)
+ batch3_data="a,104\nc,105\ne,106\n", # a:475, c:135, e:106 (new)
+ verify_initial=verify_initial,
+ verify_after_increase=verify_after_increase,
+ verify_after_decrease=verify_after_decrease,
+ )
+
+ def test_repartition_with_streaming_tws(self):
+ """Test repartition for streaming transformWithState"""
+
+ self._run_tws_repartition_test(is_pandas=False)
+
+ @unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message,
+ )
+ def test_repartition_with_streaming_tws_in_pandas(self):
+ """Test repartition for streaming transformWithStateInPandas."""
+
+ self._run_tws_repartition_test(is_pandas=True)
+
+ def _run_repartition_with_streaming_tws_multiple_state_vars_test(self, is_pandas):
+ """Test repartition with processor using multiple state variable types (value + list + map)."""
+
+ def create_streaming_df(df):
+ # Parse text input format "id,temperature" into structured columns
+ split_df = split(df["value"], ",")
+ parsed_df = df.select(
+ split_df.getItem(0).alias("id"),
+ split_df.getItem(1).cast("integer").alias("temperature"),
+ )
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value_arr", StringType(), True),
+ StructField("list_state_arr", StringType(), True),
+ StructField("map_state_arr", StringType(), True),
+ StructField("nested_map_state_arr", StringType(), True),
+ ]
+ )
+ processor = (
+ StatefulProcessorCompositeTypeFactory().pandas()
+ if is_pandas
+ else StatefulProcessorCompositeTypeFactory().row()
+ )
+ return (
+ parsed_df.groupBy("id").transformWithStateInPandas(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ if is_pandas
+ else parsed_df.groupBy("id").transformWithState(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ )
+
+ def verify_initial(results):
+ rows = {row.id: row for row in results}
+ # StatefulProcessorCompositeType initializes to "0" on first batch
+ # (this is how the processor is designed - see line 1853 in helper file)
+ self.assertEqual(rows["a"].value_arr, "0")
+ self.assertEqual(rows["a"].list_state_arr, "0")
+ # Map state initialized with default ATTRIBUTES_MAP and CONFS_MAP
+ self.assertEqual(rows["a"].map_state_arr, '{"key1": [1], "key2": [10]}')
+ self.assertEqual(rows["a"].nested_map_state_arr, '{"e1": {"e2": 5, "e3": 10}}')
+
+ def verify_after_increase(results):
+ # After repartition, state should be preserved
+ rows = {row.id: row for row in results}
+ # Key 'a': value state accumulated
+ self.assertEqual(rows["a"].value_arr, "100")
+ self.assertEqual(rows["a"].list_state_arr, "0,100")
+ # Map state updated with key "a" and temperature 100
+ self.assertEqual(rows["a"].map_state_arr, '{"a": [100], "key1": [1], "key2": [10]}')
+ self.assertEqual(
+ rows["a"].nested_map_state_arr, '{"e1": {"a": 100, "e2": 5, "e3": 10}}'
+ )
+ # New key 'd' - first batch for this key, outputs "0"
+ self.assertEqual(rows["d"].value_arr, "0")
+ self.assertEqual(rows["d"].list_state_arr, "0")
+ # Map state for 'd' initialized with defaults
+ self.assertEqual(rows["d"].map_state_arr, '{"key1": [1], "key2": [10]}')
+ self.assertEqual(rows["d"].nested_map_state_arr, '{"e1": {"e2": 5, "e3": 10}}')
+
+ def verify_after_decrease(results):
+ # After another repartition, state should still be preserved
+ rows = {row.id: row for row in results}
+ # Key 'd' gets 102, state was [0], so accumulated: [0+102] = [102]
+ self.assertEqual(rows["d"].value_arr, "102")
+ self.assertEqual(rows["d"].list_state_arr, "0,102")
+ # Map state for 'd' updated with temperature 102
+ self.assertEqual(rows["d"].map_state_arr, '{"d": [102], "key1": [1], "key2": [10]}')
+ self.assertEqual(
+ rows["d"].nested_map_state_arr, '{"e1": {"d": 102, "e2": 5, "e3": 10}}'
+ )
+ # Key 'a' gets 101, state was [100], so accumulated: [100+101] = [201]
+ self.assertEqual(rows["a"].value_arr, "201")
+ self.assertEqual(rows["a"].list_state_arr, "0,100,101")
+ # Map state for 'a' updated with temperature 101 (replaces previous value)
+ self.assertEqual(rows["a"].map_state_arr, '{"a": [101], "key1": [1], "key2": [10]}')
+ self.assertEqual(
+ rows["a"].nested_map_state_arr, '{"e1": {"a": 101, "e2": 5, "e3": 10}}'
+ )
+
+ OfflineStateRepartitionTestUtils.run_repartition_test(
+ spark=self.spark,
+ num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
+ create_streaming_df=create_streaming_df,
+ output_mode="update",
+ batch1_data="a,100\n", # a:0
+ batch2_data="a,100\nd,100\n", # a:100, d:0 (new)
+ batch3_data="d,102\na,101\n", # d:102, a:201 (new)
+ verify_initial=verify_initial,
+ verify_after_increase=verify_after_increase,
+ verify_after_decrease=verify_after_decrease,
+ )
+
+ def test_repartition_with_streaming_tws_multiple_state_vars(self):
+ """Test repartition with transformWithState using multiple state variable types (value + list + map)."""
+
+ self._run_repartition_with_streaming_tws_multiple_state_vars_test(is_pandas=False)
+
+ @unittest.skipIf(
+ not have_pandas or not have_pyarrow,
Review Comment:
ditto
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -299,6 +299,13 @@ class StatePartitionAllColumnFamiliesReader(
isTWSOperator(operatorName) && colFamilyName == StateStore.DEFAULT_COL_FAMILY_NAME
}
+
Review Comment:
remove extra line?
##########
python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py:
##########
@@ -320,6 +326,212 @@ def verify_after_decrease(results):
verify_after_decrease=verify_after_decrease,
)
+ def _run_tws_repartition_test(self, is_pandas):
+ """Helper method to run repartition test with a given processor factory method"""
+
+ def create_streaming_df(df):
+ # Parse text input format "id,temperature" into structured columns
+ split_df = split(df["value"], ",")
+ parsed_df = df.select(
+ split_df.getItem(0).alias("id"),
+ split_df.getItem(1).cast("integer").alias("temperature"),
+ )
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value", StringType(), True),
+ ]
+ )
+ processor = (
+ SimpleStatefulProcessorWithInitialStateFactory().pandas()
+ if is_pandas
+ else SimpleStatefulProcessorWithInitialStateFactory().row()
+ )
+ return (
+ parsed_df.groupBy("id").transformWithStateInPandas(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ if is_pandas
+ else parsed_df.groupBy("id").transformWithState(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ )
+
+ def verify_initial(results):
+ # SimpleStatefulProcessorWithInitialState accumulates temperature values
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "270", "a: 120 + 150 = 270")
+ self.assertEqual(values.get("b"), "50", "b: 50")
+ self.assertEqual(values.get("c"), "30", "c: 30")
+
+ def verify_after_increase(results):
+ # After repartition, state should be preserved and accumulated
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "371", "State for 'a': 270 + 101 = 371")
+ self.assertEqual(values.get("b"), "152", "State for 'b': 50 + 102 = 152")
+ self.assertEqual(values.get("d"), "103", "New key 'd' should have value 103")
+
+ def verify_after_decrease(results):
+ # After repartition, state should still be preserved
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "475", "State for 'a': 371 + 104 = 475")
+ self.assertEqual(values.get("c"), "135", "State for 'c': 30 + 105 = 135")
+ self.assertEqual(values.get("e"), "106", "New key 'e' should have value 106")
+
+ OfflineStateRepartitionTestUtils.run_repartition_test(
+ spark=self.spark,
+ num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
+ create_streaming_df=create_streaming_df,
+ output_mode="update",
+ batch1_data="a,120\na,150\nb,50\nc,30\n", # a:270, b:50, c:30
+ batch2_data="a,101\nb,102\nd,103\n", # a:371, b:152, d:103 (new)
+ batch3_data="a,104\nc,105\ne,106\n", # a:475, c:135, e:106 (new)
+ verify_initial=verify_initial,
+ verify_after_increase=verify_after_increase,
+ verify_after_decrease=verify_after_decrease,
+ )
+
+ def test_repartition_with_streaming_tws(self):
Review Comment:
for tws pandas=False tests, we add this:
https://github.com/apache/spark/blob/0cba1c0a81a8379051d009c69f618ea635485a35/python/pyspark/sql/tests/pandas/streaming/test_transform_with_state_state_variable.py#L33
##########
python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py:
##########
@@ -320,6 +326,212 @@ def verify_after_decrease(results):
verify_after_decrease=verify_after_decrease,
)
+ def _run_tws_repartition_test(self, is_pandas):
+ """Helper method to run repartition test with a given processor factory method"""
+
+ def create_streaming_df(df):
+ # Parse text input format "id,temperature" into structured columns
+ split_df = split(df["value"], ",")
+ parsed_df = df.select(
+ split_df.getItem(0).alias("id"),
+ split_df.getItem(1).cast("integer").alias("temperature"),
+ )
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value", StringType(), True),
+ ]
+ )
+ processor = (
+ SimpleStatefulProcessorWithInitialStateFactory().pandas()
+ if is_pandas
+ else SimpleStatefulProcessorWithInitialStateFactory().row()
+ )
+ return (
+ parsed_df.groupBy("id").transformWithStateInPandas(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ if is_pandas
+ else parsed_df.groupBy("id").transformWithState(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ )
+
+ def verify_initial(results):
+ # SimpleStatefulProcessorWithInitialState accumulates temperature values
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "270", "a: 120 + 150 = 270")
+ self.assertEqual(values.get("b"), "50", "b: 50")
+ self.assertEqual(values.get("c"), "30", "c: 30")
+
+ def verify_after_increase(results):
+ # After repartition, state should be preserved and accumulated
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "371", "State for 'a': 270 + 101 = 371")
+ self.assertEqual(values.get("b"), "152", "State for 'b': 50 + 102 = 152")
+ self.assertEqual(values.get("d"), "103", "New key 'd' should have value 103")
+
+ def verify_after_decrease(results):
+ # After repartition, state should still be preserved
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "475", "State for 'a': 371 + 104 = 475")
+ self.assertEqual(values.get("c"), "135", "State for 'c': 30 + 105 = 135")
+ self.assertEqual(values.get("e"), "106", "New key 'e' should have value 106")
+
+ OfflineStateRepartitionTestUtils.run_repartition_test(
+ spark=self.spark,
+ num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
+ create_streaming_df=create_streaming_df,
+ output_mode="update",
+ batch1_data="a,120\na,150\nb,50\nc,30\n", # a:270, b:50, c:30
+ batch2_data="a,101\nb,102\nd,103\n", # a:371, b:152, d:103 (new)
+ batch3_data="a,104\nc,105\ne,106\n", # a:475, c:135, e:106 (new)
+ verify_initial=verify_initial,
+ verify_after_increase=verify_after_increase,
+ verify_after_decrease=verify_after_decrease,
+ )
+
+ def test_repartition_with_streaming_tws(self):
+ """Test repartition for streaming transformWithState"""
+
+ self._run_tws_repartition_test(is_pandas=False)
+
+ @unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message,
+ )
+ def test_repartition_with_streaming_tws_in_pandas(self):
+ """Test repartition for streaming transformWithStateInPandas."""
+
+ self._run_tws_repartition_test(is_pandas=True)
+
+ def _run_repartition_with_streaming_tws_multiple_state_vars_test(self, is_pandas):
+ """Test repartition with processor using multiple state variable types (value + list + map)."""
+
+ def create_streaming_df(df):
+ # Parse text input format "id,temperature" into structured columns
+ split_df = split(df["value"], ",")
+ parsed_df = df.select(
+ split_df.getItem(0).alias("id"),
+ split_df.getItem(1).cast("integer").alias("temperature"),
+ )
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value_arr", StringType(), True),
+ StructField("list_state_arr", StringType(), True),
+ StructField("map_state_arr", StringType(), True),
+ StructField("nested_map_state_arr", StringType(), True),
+ ]
+ )
+ processor = (
+ StatefulProcessorCompositeTypeFactory().pandas()
+ if is_pandas
+ else StatefulProcessorCompositeTypeFactory().row()
+ )
+ return (
+ parsed_df.groupBy("id").transformWithStateInPandas(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ if is_pandas
+ else parsed_df.groupBy("id").transformWithState(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ )
+
+ def verify_initial(results):
+ rows = {row.id: row for row in results}
+ # StatefulProcessorCompositeType initializes to "0" on first batch
+ # (this is how the processor is designed - see line 1853 in helper file)
+ self.assertEqual(rows["a"].value_arr, "0")
+ self.assertEqual(rows["a"].list_state_arr, "0")
+ # Map state initialized with default ATTRIBUTES_MAP and CONFS_MAP
+ self.assertEqual(rows["a"].map_state_arr, '{"key1": [1], "key2": [10]}')
+ self.assertEqual(rows["a"].nested_map_state_arr, '{"e1": {"e2": 5, "e3": 10}}')
+
+ def verify_after_increase(results):
+ # After repartition, state should be preserved
+ rows = {row.id: row for row in results}
+ # Key 'a': value state accumulated
+ self.assertEqual(rows["a"].value_arr, "100")
+ self.assertEqual(rows["a"].list_state_arr, "0,100")
+ # Map state updated with key "a" and temperature 100
+ self.assertEqual(rows["a"].map_state_arr, '{"a": [100], "key1": [1], "key2": [10]}')
+ self.assertEqual(
+ rows["a"].nested_map_state_arr, '{"e1": {"a": 100, "e2": 5, "e3": 10}}'
+ )
+ # New key 'd' - first batch for this key, outputs "0"
+ self.assertEqual(rows["d"].value_arr, "0")
+ self.assertEqual(rows["d"].list_state_arr, "0")
+ # Map state for 'd' initialized with defaults
+ self.assertEqual(rows["d"].map_state_arr, '{"key1": [1], "key2": [10]}')
+ self.assertEqual(rows["d"].nested_map_state_arr, '{"e1": {"e2": 5, "e3": 10}}')
+
+ def verify_after_decrease(results):
+ # After another repartition, state should still be preserved
+ rows = {row.id: row for row in results}
+ # Key 'd' gets 102, state was [0], so accumulated: [0+102] = [102]
+ self.assertEqual(rows["d"].value_arr, "102")
+ self.assertEqual(rows["d"].list_state_arr, "0,102")
+ # Map state for 'd' updated with temperature 102
+ self.assertEqual(rows["d"].map_state_arr, '{"d": [102], "key1": [1], "key2": [10]}')
+ self.assertEqual(
+ rows["d"].nested_map_state_arr, '{"e1": {"d": 102, "e2": 5, "e3": 10}}'
+ )
+ # Key 'a' gets 101, state was [100], so accumulated: [100+101] = [201]
+ self.assertEqual(rows["a"].value_arr, "201")
+ self.assertEqual(rows["a"].list_state_arr, "0,100,101")
+ # Map state for 'a' updated with temperature 101 (replaces previous value)
+ self.assertEqual(rows["a"].map_state_arr, '{"a": [101], "key1": [1], "key2": [10]}')
+ self.assertEqual(
+ rows["a"].nested_map_state_arr, '{"e1": {"a": 101, "e2": 5, "e3": 10}}'
+ )
+
+ OfflineStateRepartitionTestUtils.run_repartition_test(
+ spark=self.spark,
+ num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
+ create_streaming_df=create_streaming_df,
+ output_mode="update",
+ batch1_data="a,100\n", # a:0
+ batch2_data="a,100\nd,100\n", # a:100, d:0 (new)
+ batch3_data="d,102\na,101\n", # d:102, a:201 (new)
+ verify_initial=verify_initial,
+ verify_after_increase=verify_after_increase,
+ verify_after_decrease=verify_after_decrease,
+ )
+
+ def test_repartition_with_streaming_tws_multiple_state_vars(self):
Review Comment:
ditto
--
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-55246][SS] Add Test for Pyspark TWS and TWSInPandas and Fix StatePartitionAllColumnFamiliesWriter Bug [spark]
Posted by "anishshri-db (via GitHub)" <gi...@apache.org>.
anishshri-db closed pull request #54019: [SPARK-55246][SS] Add Test for Pyspark TWS and TWSInPandas and Fix StatePartitionAllColumnFamiliesWriter Bug
URL: https://github.com/apache/spark/pull/54019
--
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-55246][SS] Add Test for Pyspark TWS and TWSInPandas and Fix StatePartitionAllColumnFamiliesWriter Bug [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on code in PR #54019:
URL: https://github.com/apache/spark/pull/54019#discussion_r2735039854
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StatePartitionWriter.scala:
##########
@@ -55,24 +53,16 @@ class StatePartitionAllColumnFamiliesWriter(
storeName: String,
currentBatchId: Long,
colFamilyToWriterInfoMap: Map[String, StatePartitionWriterColumnFamilyInfo],
- operatorName: String,
- schemaProviderOpt: Option[StateSchemaProvider],
- sqlConf: SQLConf) {
-
- private def isJoinV3Operator(
- operatorName: String, sqlConf: SQLConf): Boolean = {
- operatorName == StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME &&
- sqlConf.getConf(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION) == 3
- }
+ schemaProviderOpt: Option[StateSchemaProvider]) {
+ // Using the heuristic that all operators that enable column families
+ // has a non-default column family
+ private val useColumnFamilies = colFamilyToWriterInfoMap.keys.toSeq
+ .exists(_ != StateStoreId.DEFAULT_STORE_NAME)
private val defaultSchema = {
colFamilyToWriterInfoMap.get(StateStore.DEFAULT_COL_FAMILY_NAME) match {
case Some(info) => info.schema
case None =>
- // joinV3 operator doesn't have default column family schema
- assert(isJoinV3Operator(operatorName, sqlConf),
Review Comment:
fyi @zifeif2 so you don't miss this 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-55246][SS] Add Test for Pyspark TWS and TWSInPandas and Fix StatePartitionAllColumnFamiliesWriter Bug [spark]
Posted by "micheal-o (via GitHub)" <gi...@apache.org>.
micheal-o commented on code in PR #54019:
URL: https://github.com/apache/spark/pull/54019#discussion_r2734919596
##########
python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py:
##########
@@ -320,6 +324,87 @@ def verify_after_decrease(results):
verify_after_decrease=verify_after_decrease,
)
+ def _run_tws_repartition_test(self, processor_factory, is_pandas):
+ """Helper method to run repartition test with a given processor factory method"""
+ from pyspark.sql.functions import split
+
+ def create_streaming_df(df):
+ # Parse text input format "id,temperature" into structured columns
+ split_df = split(df["value"], ",")
+ parsed_df = df.select(
+ split_df.getItem(0).alias("id"),
+ split_df.getItem(1).cast("integer").alias("temperature"),
+ )
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value", StringType(), True),
+ ]
+ )
+ return (
+ parsed_df.groupBy("id").transformWithStateInPandas(
+ statefulProcessor=processor_factory.pandas(),
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ if is_pandas
+ else parsed_df.groupBy("id").transformWithState(
+ statefulProcessor=processor_factory.row(),
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ )
+
+ def verify_initial(results):
+ # SimpleStatefulProcessorWithInitialState accumulates temperature values
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "270", "a: 120 + 150 = 270")
+ self.assertEqual(values.get("b"), "50", "b: 50")
+ self.assertEqual(values.get("c"), "30", "c: 30")
+
+ def verify_after_increase(results):
+ # After repartition, state should be preserved and accumulated
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "371", "State for 'a': 270 + 101 = 371")
+ self.assertEqual(values.get("b"), "152", "State for 'b': 50 + 102 = 152")
+ self.assertEqual(values.get("d"), "103", "New key 'd' should have value 103")
+
+ def verify_after_decrease(results):
+ # After repartition, state should still be preserved
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "475", "State for 'a': 371 + 104 = 475")
+ self.assertEqual(values.get("c"), "135", "State for 'c': 30 + 105 = 135")
+ self.assertEqual(values.get("e"), "106", "New key 'e' should have value 106")
+
+ OfflineStateRepartitionTestUtils.run_repartition_test(
+ spark=self.spark,
+ num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
+ create_streaming_df=create_streaming_df,
+ output_mode="update",
+ batch1_data="a,120\na,150\nb,50\nc,30\n", # a:270, b:50, c:30
+ batch2_data="a,101\nb,102\nd,103\n", # a:371, b:152, d:103 (new)
+ batch3_data="a,104\nc,105\ne,106\n", # a:475, c:135, e:106 (new)
+ verify_initial=verify_initial,
+ verify_after_increase=verify_after_increase,
+ verify_after_decrease=verify_after_decrease,
+ )
+
+ def test_repartition_with_streaming_tws(self):
+ """Test repartition for streaming transformWithState with both row and pandas processors"""
Review Comment:
fix comment
##########
python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py:
##########
@@ -320,6 +324,87 @@ def verify_after_decrease(results):
verify_after_decrease=verify_after_decrease,
)
+ def _run_tws_repartition_test(self, processor_factory, is_pandas):
+ """Helper method to run repartition test with a given processor factory method"""
+ from pyspark.sql.functions import split
+
+ def create_streaming_df(df):
+ # Parse text input format "id,temperature" into structured columns
+ split_df = split(df["value"], ",")
+ parsed_df = df.select(
+ split_df.getItem(0).alias("id"),
+ split_df.getItem(1).cast("integer").alias("temperature"),
+ )
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value", StringType(), True),
+ ]
+ )
+ return (
+ parsed_df.groupBy("id").transformWithStateInPandas(
+ statefulProcessor=processor_factory.pandas(),
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ if is_pandas
+ else parsed_df.groupBy("id").transformWithState(
+ statefulProcessor=processor_factory.row(),
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ )
+
+ def verify_initial(results):
+ # SimpleStatefulProcessorWithInitialState accumulates temperature values
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "270", "a: 120 + 150 = 270")
+ self.assertEqual(values.get("b"), "50", "b: 50")
+ self.assertEqual(values.get("c"), "30", "c: 30")
+
+ def verify_after_increase(results):
+ # After repartition, state should be preserved and accumulated
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "371", "State for 'a': 270 + 101 = 371")
+ self.assertEqual(values.get("b"), "152", "State for 'b': 50 + 102 = 152")
+ self.assertEqual(values.get("d"), "103", "New key 'd' should have value 103")
+
+ def verify_after_decrease(results):
+ # After repartition, state should still be preserved
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "475", "State for 'a': 371 + 104 = 475")
+ self.assertEqual(values.get("c"), "135", "State for 'c': 30 + 105 = 135")
+ self.assertEqual(values.get("e"), "106", "New key 'e' should have value 106")
+
+ OfflineStateRepartitionTestUtils.run_repartition_test(
+ spark=self.spark,
+ num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
+ create_streaming_df=create_streaming_df,
+ output_mode="update",
+ batch1_data="a,120\na,150\nb,50\nc,30\n", # a:270, b:50, c:30
+ batch2_data="a,101\nb,102\nd,103\n", # a:371, b:152, d:103 (new)
+ batch3_data="a,104\nc,105\ne,106\n", # a:475, c:135, e:106 (new)
+ verify_initial=verify_initial,
+ verify_after_increase=verify_after_increase,
+ verify_after_decrease=verify_after_decrease,
+ )
+
+ def test_repartition_with_streaming_tws(self):
+ """Test repartition for streaming transformWithState with both row and pandas processors"""
+ self._run_tws_repartition_test(
+ SimpleStatefulProcessorWithInitialStateFactory(), is_pandas=False
Review Comment:
We should also add a test case that uses a processor with multiple state vars (i.e. value + list + map). Just to make sure there is coverage. There should be an existing processor for this.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StatePartitionWriter.scala:
##########
@@ -55,24 +53,16 @@ class StatePartitionAllColumnFamiliesWriter(
storeName: String,
currentBatchId: Long,
colFamilyToWriterInfoMap: Map[String, StatePartitionWriterColumnFamilyInfo],
- operatorName: String,
- schemaProviderOpt: Option[StateSchemaProvider],
- sqlConf: SQLConf) {
-
- private def isJoinV3Operator(
- operatorName: String, sqlConf: SQLConf): Boolean = {
- operatorName == StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME &&
- sqlConf.getConf(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION) == 3
- }
+ schemaProviderOpt: Option[StateSchemaProvider]) {
+ // Using the heuristic that all operators that enable column families
+ // has a non-default column family
+ private val useColumnFamilies = colFamilyToWriterInfoMap.keys.toSeq
+ .exists(_ != StateStoreId.DEFAULT_STORE_NAME)
private val defaultSchema = {
colFamilyToWriterInfoMap.get(StateStore.DEFAULT_COL_FAMILY_NAME) match {
case Some(info) => info.schema
case None =>
- // joinV3 operator doesn't have default column family schema
- assert(isJoinV3Operator(operatorName, sqlConf),
Review Comment:
Lets add the assert. Makes it easier to know/debug if we run into that situation
##########
python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py:
##########
@@ -29,6 +29,10 @@
pandas_requirement_message,
pyarrow_requirement_message,
)
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ TTLStatefulProcessorFactory,
+ SimpleStatefulProcessorWithInitialStateFactory,
Review Comment:
S before T
##########
python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py:
##########
@@ -29,6 +29,10 @@
pandas_requirement_message,
pyarrow_requirement_message,
)
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
Review Comment:
move up
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StatePartitionWriter.scala:
##########
@@ -55,24 +53,16 @@ class StatePartitionAllColumnFamiliesWriter(
storeName: String,
currentBatchId: Long,
colFamilyToWriterInfoMap: Map[String, StatePartitionWriterColumnFamilyInfo],
- operatorName: String,
- schemaProviderOpt: Option[StateSchemaProvider],
- sqlConf: SQLConf) {
-
- private def isJoinV3Operator(
- operatorName: String, sqlConf: SQLConf): Boolean = {
- operatorName == StatefulOperatorsUtils.SYMMETRIC_HASH_JOIN_EXEC_OP_NAME &&
- sqlConf.getConf(SQLConf.STREAMING_JOIN_STATE_FORMAT_VERSION) == 3
- }
+ schemaProviderOpt: Option[StateSchemaProvider]) {
+ // Using the heuristic that all operators that enable column families
+ // have a non-default column family
+ private val useColumnFamilies = colFamilyToWriterInfoMap.keys.toSeq
+ .exists(_ != StateStoreId.DEFAULT_STORE_NAME)
Review Comment:
`StateStore.DEFAULT_COL_FAMILY_NAME`
##########
python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py:
##########
@@ -29,6 +29,10 @@
pandas_requirement_message,
Review Comment:
run `dev/reformat-python` to fix formatting issue
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -299,6 +299,13 @@ class StatePartitionAllColumnFamiliesReader(
isTWSOperator(operatorName) && colFamilyName == StateStore.DEFAULT_COL_FAMILY_NAME
}
+
+ private def checkIfUseColumnFamily(schema: Set[StateStoreColFamilySchema]): Boolean = {
+ // Using the heuristic that all operators that enable column families
+ // have a non-default column family
+ schema.exists(_.colFamilyName != StateStoreId.DEFAULT_STORE_NAME)
Review Comment:
`StateStore.DEFAULT_COL_FAMILY_NAME`
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -299,6 +299,13 @@ class StatePartitionAllColumnFamiliesReader(
isTWSOperator(operatorName) && colFamilyName == StateStore.DEFAULT_COL_FAMILY_NAME
}
+
+ private def checkIfUseColumnFamily(schema: Set[StateStoreColFamilySchema]): Boolean = {
Review Comment:
can just name it `useColumnFamilies`
##########
python/pyspark/sql/tests/streaming/test_streaming_offline_state_repartition.py:
##########
@@ -320,6 +324,87 @@ def verify_after_decrease(results):
verify_after_decrease=verify_after_decrease,
)
+ def _run_tws_repartition_test(self, processor_factory, is_pandas):
+ """Helper method to run repartition test with a given processor factory method"""
+ from pyspark.sql.functions import split
+
+ def create_streaming_df(df):
+ # Parse text input format "id,temperature" into structured columns
+ split_df = split(df["value"], ",")
+ parsed_df = df.select(
+ split_df.getItem(0).alias("id"),
+ split_df.getItem(1).cast("integer").alias("temperature"),
+ )
+
+ output_schema = StructType(
+ [
+ StructField("id", StringType(), True),
+ StructField("value", StringType(), True),
+ ]
+ )
+ return (
+ parsed_df.groupBy("id").transformWithStateInPandas(
+ statefulProcessor=processor_factory.pandas(),
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ if is_pandas
+ else parsed_df.groupBy("id").transformWithState(
+ statefulProcessor=processor_factory.row(),
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ initialState=None,
+ )
+ )
+
+ def verify_initial(results):
+ # SimpleStatefulProcessorWithInitialState accumulates temperature values
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "270", "a: 120 + 150 = 270")
+ self.assertEqual(values.get("b"), "50", "b: 50")
+ self.assertEqual(values.get("c"), "30", "c: 30")
+
+ def verify_after_increase(results):
+ # After repartition, state should be preserved and accumulated
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "371", "State for 'a': 270 + 101 = 371")
+ self.assertEqual(values.get("b"), "152", "State for 'b': 50 + 102 = 152")
+ self.assertEqual(values.get("d"), "103", "New key 'd' should have value 103")
+
+ def verify_after_decrease(results):
+ # After repartition, state should still be preserved
+ values = {row.id: row.value for row in results}
+ self.assertEqual(values.get("a"), "475", "State for 'a': 371 + 104 = 475")
+ self.assertEqual(values.get("c"), "135", "State for 'c': 30 + 105 = 135")
+ self.assertEqual(values.get("e"), "106", "New key 'e' should have value 106")
+
+ OfflineStateRepartitionTestUtils.run_repartition_test(
+ spark=self.spark,
+ num_shuffle_partitions=self.NUM_SHUFFLE_PARTITIONS,
+ create_streaming_df=create_streaming_df,
+ output_mode="update",
+ batch1_data="a,120\na,150\nb,50\nc,30\n", # a:270, b:50, c:30
+ batch2_data="a,101\nb,102\nd,103\n", # a:371, b:152, d:103 (new)
+ batch3_data="a,104\nc,105\ne,106\n", # a:475, c:135, e:106 (new)
+ verify_initial=verify_initial,
+ verify_after_increase=verify_after_increase,
+ verify_after_decrease=verify_after_decrease,
+ )
+
+ def test_repartition_with_streaming_tws(self):
+ """Test repartition for streaming transformWithState with both row and pandas processors"""
+ self._run_tws_repartition_test(
+ SimpleStatefulProcessorWithInitialStateFactory(), is_pandas=False
+ )
+
+ def test_repartition_with_streaming_tws_in_pandas(self):
Review Comment:
pandas tests require: https://github.com/apache/spark/blob/9ce067c7d532a6766d622e55164baa942a065250/python/pyspark/sql/tests/pandas/streaming/test_pandas_transform_with_state_state_variable.py#L979
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/state/StatePartitionReader.scala:
##########
@@ -299,6 +299,13 @@ class StatePartitionAllColumnFamiliesReader(
isTWSOperator(operatorName) && colFamilyName == StateStore.DEFAULT_COL_FAMILY_NAME
}
+
+ private def checkIfUseColumnFamily(schema: Set[StateStoreColFamilySchema]): Boolean = {
Review Comment:
This should be a `private lazy val` right? We don't need to compute this every time right? since the `stateStoreColFamilySchemas` isn't changing within the class
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55246][SS] Add Test for Pyspark TWS and TWSInPandas and Fix StatePartitionAllColumnFamiliesWriter Bug [spark]
Posted by "zifeif2 (via GitHub)" <gi...@apache.org>.
zifeif2 commented on code in PR #54019:
URL: https://github.com/apache/spark/pull/54019#discussion_r2738897833
##########
python/pyspark/sql/tests/pandas/helper/helper_pandas_transform_with_state.py:
##########
@@ -724,8 +724,8 @@ def handleInputRows(self, key, rows, timerValues) -> Iterator[pd.DataFrame]:
if self.ttl_count_state.exists():
ttl_count = self.ttl_count_state.get()[0]
if self.ttl_list_state.exists():
- iter = self.ttl_list_state.get()
- for s in iter:
+ list_iter = self.ttl_list_state.get()
Review Comment:
These processors are not used in this test, but still renaming the variable so that it doesn't have a name collision with "iter" function.
--
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