You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "fedimser (via GitHub)" <gi...@apache.org> on 2026/01/10 04:24:22 UTC
[PR] [SPARK-54805] Implement TwsTester in PySpark [spark]
fedimser opened a new pull request, #53758:
URL: https://github.com/apache/spark/pull/53758
<!--
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.
-->
* Added TwsTester, a test helper for writing unit tests for StatefulProcessor implementations that will be used in TrwansformWithState operator in streaming queries. It processes input rows and returns output rows equivalent to those that would be produced by the processor in an actual Spark streaming query.
* Supported functionality:
* Processing input rows and producing output rows via `test()`.
* Initial state setup via constructor parameter.
* Direct state manipulation via `setValueState`, `setListState`, `setMapState`.
* Direct state inspection via `peekValueState`, `peekListState`, `peekMapState`.
* Timers (both ProcessingTime and EventTime modes).
* TTL.
* Real-time mode simulation.
* This is Python version of Scala TwsTester in https://github.com/apache/spark/pull/53159.
### 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.
-->
Some users requested unit testing functionality for TWS.
### 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'.
-->
Yes, it adds new public API to PySpark:
* pyspark.sql.streaming.TwsTester
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
Added unit and end-to-end tests in this PR. End-to-end tests compare TwsTester output with results of a real streaming query.
### 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. Cursor with claude-4.5-opus-high was used to assist with coding and generate some of documentation and tests.
Generated-by: claude-4.5-opus-high
--
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR closed pull request #53758: [SPARK-54805][SS][PYTHON] Implement TwsTester in PySpark
URL: https://github.com/apache/spark/pull/53758
--
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r3308822731
##########
python/pyspark/sql/streaming/__init__.py:
##########
@@ -23,4 +23,5 @@
StatefulProcessor, # noqa: F401
StatefulProcessorHandle, # noqa: F401
) # noqa: F401
+from pyspark.sql.streaming.tws_tester import TwsTester # noqa: F401
Review Comment:
@fedimser @HeartSaVioR is this a new feature in pyspark.sql.streaming?
do we want to add it into the 4.2 API references?
cc @huaxingao @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-54805] Implement TwsTester in PySpark [spark]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #53758:
URL: https://github.com/apache/spark/pull/53758#issuecomment-3731806588
## JIRA Issue Information
=== New Feature SPARK-54805 ===
Summary: Improve the testing experience for TransformWithState in PySpark
Assignee: None
Status: Open
Affected: ["4.2.0"]
---
*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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #53758:
URL: https://github.com/apache/spark/pull/53758#issuecomment-3850579556
Probably yes if you don't change the code but to retrigger the CI.
--
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on PR #53758:
URL: https://github.com/apache/spark/pull/53758#issuecomment-3850320804
Tests that were failing now pass. There is only pyspark-connect test failing, timing out (twice). I think it is unrelated to this change, but let's try to get the green check. I rebased again to re-trigger all tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r2748042109
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ initial_state_call_count = [0]
+
+ class InitialStateCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ initial_state_call_count[0] += 1
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitialStateCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(initial_state_call_count[0], 0)
+
+ def test_delete_value_state(self):
+ """Test that deleteState correctly deletes value state for a given key."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateValueState("count", "key1", (10,))
+ tester.updateValueState("count", "key2", (20,))
+ self.assertEqual(tester.peekValueState("count", "key1"), (10,))
+
+ tester.deleteState("count", "key1")
+ self.assertIsNone(tester.peekValueState("count", "key1"))
+ self.assertEqual(tester.peekValueState("count", "key2"), (20,))
+
+ def test_delete_list_state(self):
+ """Test that deleteState correctly deletes list state for a given key."""
+ processor = TopKProcessorFactory(k=3).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
+ tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])
+
+ tester.deleteState("topK", "key1")
+ self.assertEqual(tester.peekListState("topK", "key1"), [])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])
+
+ def test_delete_map_state(self):
+ """Test that deleteState correctly deletes map state for a given key."""
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (5,), ("world",): (3,)},
+ )
+
+ tester.deleteState("frequencies", "user1")
+ self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})
+
+ def test_delete_nonexistent_state_raises_error(self):
+ """Test that deleteState raises an error for non-existent state."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
+ tester.deleteState("nonexistent_state", "key1")
+
+ # Timer tests.
+
+ def test_processing_time_timers(self):
+ """Test that TwsTester supports ProcessingTime timers."""
+ processor = SessionTimeoutProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.test("key1", [Row(value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello")])
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(expired1, [])
+
+ # Process input for key2 at t=5000 - should register timer at t=15000.
+ result2 = tester.test("key2", [Row(value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world")])
+
+ # Set processing time to 11000 - key1's timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ self.assertEqual(expired2, [Row(key="key1", result="session-expired")])
+
+ # Set processing time to 16000 - key2's timer should fire.
+ expired3 = tester.setProcessingTime(16000)
+ self.assertEqual(expired3, [Row(key="key2", result="session-expired")])
+
+ # Verify state is cleared after session expiry.
+ self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
+ self.assertIsNone(tester.peekValueState("lastSeen", "key2"))
+
+ def test_processing_time_timers_pandas(self):
+ """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
+ processor = SessionTimeoutProcessorFactory().pandas()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
+ pdt.assert_frame_equal(result1, expected1, check_like=True)
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(len(expired1), 0)
+
+ # Set processing time to 11000 - timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ expected2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
+ pdt.assert_frame_equal(expired2, expected2, check_like=True)
+
+ def test_event_time_timers_manual_watermark(self):
+ """Test that TwsTester supports EventTime timers fired by manual watermark advance."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process event at t=10000 for key1 - registers timer at t=15000.
+ result1 = tester.test("key1", [Row(event_time_ms=10000, value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello@10000")])
+
+ # Process event at t=11000 for key2 - registers timer at t=16000.
+ result2 = tester.test("key2", [Row(event_time_ms=11000, value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world@11000")])
+
+ # Set watermark to 15000 - key1's timer should fire.
+ expired1 = tester.setWatermark(15000)
+ self.assertEqual(len(expired1), 1)
+ self.assertEqual(expired1[0], Row(key="key1", result="session-expired@watermark=15000"))
+
+ # Set watermark to 16000 - key2's timer should fire.
+ expired2 = tester.setWatermark(16000)
+ self.assertEqual(len(expired2), 1)
+ self.assertEqual(expired2[0], Row(key="key2", result="session-expired@watermark=16000"))
+
+ # Verify state is cleared.
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key1"))
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key2"))
+
+ def test_late_event_filtering(self):
+ """Test that TwsTester filters late events in EventTime mode."""
+ processor = EventTimeCountProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Initially watermark is 0, so all events should be processed.
+ result1 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=5000, value="a"),
+ Row(event_time_ms=6000, value="b"),
+ Row(event_time_ms=7000, value="c"),
+ ],
+ )
+ self.assertEqual(result1, [Row(key="key1", count=3)])
+ self.assertEqual(tester.peekValueState("count", "key1"), (3,))
+
+ # Set watermark to 6000.
+ tester.setWatermark(6000)
+
+ # Now watermark is 6000. Send events with mixed event times:
+ # - (4000, "late1") -> event time 4000 <= watermark 6000, should be FILTERED.
+ # - (5000, "late2") -> event time 5000 <= watermark 6000, should be FILTERED.
+ # - (6000, "late3") -> event time 6000 <= watermark 6000, should be FILTERED.
+ # - (8000, "ontime") -> event time 8000 > watermark 6000, should be PROCESSED.
+ result2 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=4000, value="late1"),
+ Row(event_time_ms=5000, value="late2"),
+ Row(event_time_ms=6000, value="late3"),
+ Row(event_time_ms=8000, value="ontime"),
+ ],
+ )
+ # Only 1 event should be processed (the on-time one)
+ self.assertEqual(result2, [Row(key="key1", count=4)]) # 3 + 1 = 4
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+
+ def test_all_late_events_filtered(self):
+ """Test that all late events are filtered when all are older than watermark."""
+ processor = EventTimeCountProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process events to get initial count.
+ tester.test("key1", [Row(event_time_ms=10000, value="a")])
+ self.assertEqual(tester.peekValueState("count", "key1"), (1,))
+
+ # Set watermark to 10000.
+ tester.setWatermark(10000)
+
+ # Now send only late events - all should be filtered, count unchanged.
+ result = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=1000, value="late1"),
+ Row(event_time_ms=5000, value="late2"),
+ Row(event_time_ms=9999, value="late3"),
+ Row(event_time_ms=10000, value="late4"), # exactly at watermark, also filtered
+ ],
+ )
+ # No events processed - handleInputRows is called with empty list.
+ self.assertEqual(result, [Row(key="key1", count=1)]) # Count still 1
+ self.assertEqual(tester.peekValueState("count", "key1"), (1,))
+
+ def test_timer_registration_in_none_mode_raises_error(self):
+ """Test that timer operations raise error in TimeMode.None."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_set_processing_time_in_wrong_mode_raises_error(self):
+ """Test that setProcessingTime raises error in non-ProcessingTime mode."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ # Test in None mode.
+ tester_none = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_none.setProcessingTime(1000)
+
+ def test_set_watermark_in_wrong_mode_raises_error(self):
+ """Test that setWatermark raises error in non-EventTime mode."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ # Test in None mode.
+ tester_none = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_none.setWatermark(1000)
+
+ # Test in ProcessingTime mode.
+ tester_pt = TwsTester(processor, timeMode="ProcessingTime")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_pt.setWatermark(1000)
+
+ def test_eventtime_mode_requires_extractor(self):
+ """Test that EventTime mode requires eventTimeExtractor."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ with self.assertRaisesRegex(PySparkValueError, "ARGUMENT_REQUIRED"):
+ TwsTester(processor, timeMode="EventTime")
+
+ def test_processing_time_must_move_forward(self):
+ """Test that setProcessingTime requires time to move forward."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ tester.setProcessingTime(1000)
+
+ # Same time should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setProcessingTime(1000)
+
+ # Earlier time should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setProcessingTime(500)
+
+ def test_watermark_must_move_forward(self):
+ """Test that setWatermark requires watermark to move forward."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ tester.setWatermark(1000)
+
+ # Same watermark should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setWatermark(1000)
+
+ # Earlier watermark should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setWatermark(500)
+
+ def test_timer_cannot_be_registered_in_past(self):
+ """Test that timer cannot be registered at or before current watermark in EventTime mode."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Set watermark to 10000.
+ tester.setWatermark(10000)
+
+ # Try to register a timer at or before watermark - should raise error.
+ tester.handle.setGroupingKey("key1")
+ with self.assertRaises(IllegalArgumentException):
+ tester.handle.registerTimer(10000) # exactly at watermark
+
+ with self.assertRaises(IllegalArgumentException):
+ tester.handle.registerTimer(5000) # before watermark
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterFuzzTests(ReusedSQLTestCase):
+ @classmethod
+ def conf(cls):
+ cfg = SparkConf()
+ cfg.set("spark.sql.shuffle.partitions", "1")
+ cfg.set(
+ "spark.sql.streaming.stateStore.providerClass",
+ "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider",
+ )
+ return cfg
+
+ # Helper that runs real streaming query with TWS and compares results with TwsTester output.
+ # Supports both row mode (use_pandas=False) and Pandas mode (use_pandas=True).
+ # Supports multiple batches.
+ # Assumes that the first column is the key column.
+ # Output rows are checked ignoring order.
+ # Correct only if results do not depend on reordering rows within batch.
+ def _check_tws_tester(
+ self,
+ processor: StatefulProcessor,
+ batches: list[list[tuple]],
+ input_columns: list[str],
+ input_types: list[str],
+ output_schema: StructType,
+ use_pandas: bool = False,
+ ):
+ num_batches = len(batches)
+ key_column: str = input_columns[0] # Assume first column is the key column.
+ value_columns: list[str] = input_columns[1:] # All columns except the key column.
+
+ num_input_columns: int = len(input_columns)
+ assert len(input_types) == num_input_columns
+
+ # Use TwsTester to compute expected results.
+ tester = TwsTester(processor)
+
+ def _run_tester(batch: list[tuple]) -> list[Row]:
+ # Group batch by key.
+ grouped = {}
+ for row in batch:
+ key = row[0]
+ if key not in grouped:
+ grouped[key] = []
+ grouped[key].append(row[1:])
+
+ results: list[Row] = []
+ for key, values in grouped.items():
+ if use_pandas:
+ input_to_tester = pd.DataFrame(values, columns=value_columns)
+ result_df = tester.testInPandas(key, input_to_tester)
+ results.extend(self.spark.createDataFrame(result_df).collect())
+ else:
+ input_to_tester: list[Row] = [
+ Row(**dict(zip(value_columns, v))) for v in values
+ ]
+ results.extend(tester.test(key, input_to_tester))
+ return results
+
+ expected_results: list[list[Row]] = [_run_tester(batch) for batch in batches]
+
+ # Create streaming DataFrame, it will read input from files.
+ input_path = tempfile.mkdtemp()
+ df = self.spark.readStream.format("text").load(input_path)
+ df_split = df.withColumn("split_values", split(df["value"], ","))
+ df_final = df_split.select(
+ *[
+ df_split.split_values.getItem(i).alias(col_name).cast(input_types[i])
+ for i, col_name in enumerate(input_columns)
+ ]
+ )
+ self.assertTrue(df_final.isStreaming)
+
+ # Apply TransformWithState.
+ df_grouped = df_final.groupBy(key_column)
+ if use_pandas:
+ tws_df = df_grouped.transformWithStateInPandas(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ )
+ else:
+ tws_df = df_grouped.transformWithState(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ )
+
+ # Start streaming query collecting TWS results.
+ actual_results: list[list[Row]] = []
+
+ def _collect_results(batch_df: DataFrame, batch_id):
+ rows: list[Row] = batch_df.collect()
+ actual_results.append(rows)
+
+ query: StreamingQuery = (
+ tws_df.writeStream.queryName("test_tmp_query")
+ .foreachBatch(_collect_results)
+ .outputMode("update")
+ .start()
+ )
+
+ # Process each input batch.
+ def _write_batch(batch_id: int, batch: list[tuple]):
+ with open(input_path + f"/test{batch_id}.txt", "w") as f:
+ for row in batch:
+ assert len(row) == num_input_columns
+ f.write(",".join(str(value) for value in row) + "\n")
+
+ try:
+ for batch_id, batch in enumerate(batches):
+ _write_batch(batch_id, batch)
+ query.processAllAvailable()
+ query.awaitTermination(1)
+ self.assertTrue(query.exception() is None)
+ finally:
+ query.stop()
+
+ # Assert actual results are equal to expected results.
+ self.assertEqual(len(actual_results), num_batches)
+ self.assertEqual(len(expected_results), num_batches)
+ for batch_id in range(num_batches):
+ self.assertCountEqual(actual_results[batch_id], expected_results[batch_id])
Review Comment:
Wow, I just had the same impression after reading your comment and had to check again.
Here are people complaining about this name being confusing: https://discuss.python.org/t/unittest-assertcountequal-naming-is-confuse/6480
I added a comment explaining what this does.
--
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r2756134084
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ initial_state_call_count = [0]
+
+ class InitialStateCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ initial_state_call_count[0] += 1
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitialStateCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(initial_state_call_count[0], 0)
+
+ def test_delete_value_state(self):
+ """Test that deleteState correctly deletes value state for a given key."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateValueState("count", "key1", (10,))
+ tester.updateValueState("count", "key2", (20,))
+ self.assertEqual(tester.peekValueState("count", "key1"), (10,))
+
+ tester.deleteState("count", "key1")
+ self.assertIsNone(tester.peekValueState("count", "key1"))
+ self.assertEqual(tester.peekValueState("count", "key2"), (20,))
+
+ def test_delete_list_state(self):
+ """Test that deleteState correctly deletes list state for a given key."""
+ processor = TopKProcessorFactory(k=3).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
+ tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])
+
+ tester.deleteState("topK", "key1")
+ self.assertEqual(tester.peekListState("topK", "key1"), [])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])
+
+ def test_delete_map_state(self):
+ """Test that deleteState correctly deletes map state for a given key."""
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (5,), ("world",): (3,)},
+ )
+
+ tester.deleteState("frequencies", "user1")
+ self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})
+
+ def test_delete_nonexistent_state_raises_error(self):
+ """Test that deleteState raises an error for non-existent state."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
+ tester.deleteState("nonexistent_state", "key1")
+
+ # Timer tests.
+
+ def test_processing_time_timers(self):
+ """Test that TwsTester supports ProcessingTime timers."""
+ processor = SessionTimeoutProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.test("key1", [Row(value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello")])
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(expired1, [])
+
+ # Process input for key2 at t=5000 - should register timer at t=15000.
+ result2 = tester.test("key2", [Row(value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world")])
+
+ # Set processing time to 11000 - key1's timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ self.assertEqual(expired2, [Row(key="key1", result="session-expired")])
+
+ # Set processing time to 16000 - key2's timer should fire.
+ expired3 = tester.setProcessingTime(16000)
+ self.assertEqual(expired3, [Row(key="key2", result="session-expired")])
+
+ # Verify state is cleared after session expiry.
+ self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
+ self.assertIsNone(tester.peekValueState("lastSeen", "key2"))
+
+ def test_processing_time_timers_pandas(self):
+ """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
+ processor = SessionTimeoutProcessorFactory().pandas()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
+ pdt.assert_frame_equal(result1, expected1, check_like=True)
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(len(expired1), 0)
+
+ # Set processing time to 11000 - timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ expected2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
+ pdt.assert_frame_equal(expired2, expected2, check_like=True)
+
+ def test_event_time_timers_manual_watermark(self):
+ """Test that TwsTester supports EventTime timers fired by manual watermark advance."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process event at t=10000 for key1 - registers timer at t=15000.
+ result1 = tester.test("key1", [Row(event_time_ms=10000, value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello@10000")])
+
+ # Process event at t=11000 for key2 - registers timer at t=16000.
+ result2 = tester.test("key2", [Row(event_time_ms=11000, value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world@11000")])
+
+ # Set watermark to 15000 - key1's timer should fire.
+ expired1 = tester.setWatermark(15000)
+ self.assertEqual(len(expired1), 1)
+ self.assertEqual(expired1[0], Row(key="key1", result="session-expired@watermark=15000"))
+
+ # Set watermark to 16000 - key2's timer should fire.
+ expired2 = tester.setWatermark(16000)
+ self.assertEqual(len(expired2), 1)
+ self.assertEqual(expired2[0], Row(key="key2", result="session-expired@watermark=16000"))
+
+ # Verify state is cleared.
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key1"))
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key2"))
+
+ def test_late_event_filtering(self):
+ """Test that TwsTester filters late events in EventTime mode."""
+ processor = EventTimeCountProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Initially watermark is 0, so all events should be processed.
+ result1 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=5000, value="a"),
+ Row(event_time_ms=6000, value="b"),
+ Row(event_time_ms=7000, value="c"),
+ ],
+ )
+ self.assertEqual(result1, [Row(key="key1", count=3)])
+ self.assertEqual(tester.peekValueState("count", "key1"), (3,))
+
+ # Set watermark to 6000.
+ tester.setWatermark(6000)
+
+ # Now watermark is 6000. Send events with mixed event times:
+ # - (4000, "late1") -> event time 4000 <= watermark 6000, should be FILTERED.
+ # - (5000, "late2") -> event time 5000 <= watermark 6000, should be FILTERED.
+ # - (6000, "late3") -> event time 6000 <= watermark 6000, should be FILTERED.
+ # - (8000, "ontime") -> event time 8000 > watermark 6000, should be PROCESSED.
+ result2 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=4000, value="late1"),
+ Row(event_time_ms=5000, value="late2"),
+ Row(event_time_ms=6000, value="late3"),
+ Row(event_time_ms=8000, value="ontime"),
+ ],
+ )
+ # Only 1 event should be processed (the on-time one)
+ self.assertEqual(result2, [Row(key="key1", count=4)]) # 3 + 1 = 4
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+
+ def test_all_late_events_filtered(self):
+ """Test that all late events are filtered when all are older than watermark."""
+ processor = EventTimeCountProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process events to get initial count.
+ tester.test("key1", [Row(event_time_ms=10000, value="a")])
+ self.assertEqual(tester.peekValueState("count", "key1"), (1,))
+
+ # Set watermark to 10000.
+ tester.setWatermark(10000)
+
+ # Now send only late events - all should be filtered, count unchanged.
+ result = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=1000, value="late1"),
+ Row(event_time_ms=5000, value="late2"),
+ Row(event_time_ms=9999, value="late3"),
+ Row(event_time_ms=10000, value="late4"), # exactly at watermark, also filtered
+ ],
+ )
+ # No events processed - handleInputRows is called with empty list.
+ self.assertEqual(result, [Row(key="key1", count=1)]) # Count still 1
+ self.assertEqual(tester.peekValueState("count", "key1"), (1,))
+
+ def test_timer_registration_in_none_mode_raises_error(self):
+ """Test that timer operations raise error in TimeMode.None."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_set_processing_time_in_wrong_mode_raises_error(self):
+ """Test that setProcessingTime raises error in non-ProcessingTime mode."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ # Test in None mode.
+ tester_none = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_none.setProcessingTime(1000)
+
+ def test_set_watermark_in_wrong_mode_raises_error(self):
+ """Test that setWatermark raises error in non-EventTime mode."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ # Test in None mode.
+ tester_none = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_none.setWatermark(1000)
+
+ # Test in ProcessingTime mode.
+ tester_pt = TwsTester(processor, timeMode="ProcessingTime")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_pt.setWatermark(1000)
+
+ def test_eventtime_mode_requires_extractor(self):
+ """Test that EventTime mode requires eventTimeExtractor."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ with self.assertRaisesRegex(PySparkValueError, "ARGUMENT_REQUIRED"):
+ TwsTester(processor, timeMode="EventTime")
+
+ def test_processing_time_must_move_forward(self):
+ """Test that setProcessingTime requires time to move forward."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ tester.setProcessingTime(1000)
+
+ # Same time should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setProcessingTime(1000)
+
+ # Earlier time should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setProcessingTime(500)
+
+ def test_watermark_must_move_forward(self):
+ """Test that setWatermark requires watermark to move forward."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ tester.setWatermark(1000)
+
+ # Same watermark should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setWatermark(1000)
+
+ # Earlier watermark should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setWatermark(500)
+
+ def test_timer_cannot_be_registered_in_past(self):
+ """Test that timer cannot be registered at or before current watermark in EventTime mode."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Set watermark to 10000.
+ tester.setWatermark(10000)
+
+ # Try to register a timer at or before watermark - should raise error.
+ tester.handle.setGroupingKey("key1")
+ with self.assertRaises(IllegalArgumentException):
+ tester.handle.registerTimer(10000) # exactly at watermark
+
+ with self.assertRaises(IllegalArgumentException):
+ tester.handle.registerTimer(5000) # before watermark
+
Review Comment:
Added this test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r2755795979
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ initial_state_call_count = [0]
+
+ class InitialStateCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ initial_state_call_count[0] += 1
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitialStateCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(initial_state_call_count[0], 0)
+
+ def test_delete_value_state(self):
+ """Test that deleteState correctly deletes value state for a given key."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateValueState("count", "key1", (10,))
+ tester.updateValueState("count", "key2", (20,))
+ self.assertEqual(tester.peekValueState("count", "key1"), (10,))
+
+ tester.deleteState("count", "key1")
+ self.assertIsNone(tester.peekValueState("count", "key1"))
+ self.assertEqual(tester.peekValueState("count", "key2"), (20,))
+
+ def test_delete_list_state(self):
+ """Test that deleteState correctly deletes list state for a given key."""
+ processor = TopKProcessorFactory(k=3).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
+ tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])
+
+ tester.deleteState("topK", "key1")
+ self.assertEqual(tester.peekListState("topK", "key1"), [])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])
+
+ def test_delete_map_state(self):
+ """Test that deleteState correctly deletes map state for a given key."""
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (5,), ("world",): (3,)},
+ )
+
+ tester.deleteState("frequencies", "user1")
+ self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})
+
+ def test_delete_nonexistent_state_raises_error(self):
+ """Test that deleteState raises an error for non-existent state."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
+ tester.deleteState("nonexistent_state", "key1")
+
+ # Timer tests.
+
+ def test_processing_time_timers(self):
+ """Test that TwsTester supports ProcessingTime timers."""
+ processor = SessionTimeoutProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.test("key1", [Row(value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello")])
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(expired1, [])
+
+ # Process input for key2 at t=5000 - should register timer at t=15000.
+ result2 = tester.test("key2", [Row(value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world")])
+
+ # Set processing time to 11000 - key1's timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ self.assertEqual(expired2, [Row(key="key1", result="session-expired")])
+
+ # Set processing time to 16000 - key2's timer should fire.
+ expired3 = tester.setProcessingTime(16000)
+ self.assertEqual(expired3, [Row(key="key2", result="session-expired")])
+
+ # Verify state is cleared after session expiry.
+ self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
+ self.assertIsNone(tester.peekValueState("lastSeen", "key2"))
+
+ def test_processing_time_timers_pandas(self):
+ """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
+ processor = SessionTimeoutProcessorFactory().pandas()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
+ pdt.assert_frame_equal(result1, expected1, check_like=True)
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(len(expired1), 0)
+
+ # Set processing time to 11000 - timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ expected2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
+ pdt.assert_frame_equal(expired2, expected2, check_like=True)
+
+ def test_event_time_timers_manual_watermark(self):
+ """Test that TwsTester supports EventTime timers fired by manual watermark advance."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process event at t=10000 for key1 - registers timer at t=15000.
+ result1 = tester.test("key1", [Row(event_time_ms=10000, value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello@10000")])
+
+ # Process event at t=11000 for key2 - registers timer at t=16000.
+ result2 = tester.test("key2", [Row(event_time_ms=11000, value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world@11000")])
+
+ # Set watermark to 15000 - key1's timer should fire.
+ expired1 = tester.setWatermark(15000)
+ self.assertEqual(len(expired1), 1)
+ self.assertEqual(expired1[0], Row(key="key1", result="session-expired@watermark=15000"))
+
+ # Set watermark to 16000 - key2's timer should fire.
+ expired2 = tester.setWatermark(16000)
+ self.assertEqual(len(expired2), 1)
+ self.assertEqual(expired2[0], Row(key="key2", result="session-expired@watermark=16000"))
+
+ # Verify state is cleared.
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key1"))
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key2"))
+
+ def test_late_event_filtering(self):
+ """Test that TwsTester filters late events in EventTime mode."""
+ processor = EventTimeCountProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Initially watermark is 0, so all events should be processed.
+ result1 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=5000, value="a"),
+ Row(event_time_ms=6000, value="b"),
+ Row(event_time_ms=7000, value="c"),
+ ],
+ )
+ self.assertEqual(result1, [Row(key="key1", count=3)])
+ self.assertEqual(tester.peekValueState("count", "key1"), (3,))
+
+ # Set watermark to 6000.
+ tester.setWatermark(6000)
+
+ # Now watermark is 6000. Send events with mixed event times:
+ # - (4000, "late1") -> event time 4000 <= watermark 6000, should be FILTERED.
+ # - (5000, "late2") -> event time 5000 <= watermark 6000, should be FILTERED.
+ # - (6000, "late3") -> event time 6000 <= watermark 6000, should be FILTERED.
+ # - (8000, "ontime") -> event time 8000 > watermark 6000, should be PROCESSED.
+ result2 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=4000, value="late1"),
+ Row(event_time_ms=5000, value="late2"),
+ Row(event_time_ms=6000, value="late3"),
+ Row(event_time_ms=8000, value="ontime"),
+ ],
+ )
+ # Only 1 event should be processed (the on-time one)
+ self.assertEqual(result2, [Row(key="key1", count=4)]) # 3 + 1 = 4
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+
+ def test_all_late_events_filtered(self):
Review Comment:
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r2746010874
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
Review Comment:
I see we have tests for None key and value (even grouping key) - could you please double check whether this is allowed in actual implementation of TWS? Probably better to in line with the actual impl.
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
Review Comment:
nit: Not sure about what we are specifically testing here. If we were trying to have completeness of test coverage, we probably should have test for getValueState and getMapState as well. Say, test_clear_nonexistent_map_state also makes sense if this test makes sense. If it doesn't worth adding, this test probably also does not worth it.
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
Review Comment:
nit: I feel like test_multiple_XXX_states and test_mixed_state_types tests are redundant with test_all_methods_processor. Shall we remove them?
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ initial_state_call_count = [0]
+
+ class InitialStateCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ initial_state_call_count[0] += 1
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitialStateCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(initial_state_call_count[0], 0)
+
+ def test_delete_value_state(self):
+ """Test that deleteState correctly deletes value state for a given key."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateValueState("count", "key1", (10,))
+ tester.updateValueState("count", "key2", (20,))
+ self.assertEqual(tester.peekValueState("count", "key1"), (10,))
+
+ tester.deleteState("count", "key1")
+ self.assertIsNone(tester.peekValueState("count", "key1"))
+ self.assertEqual(tester.peekValueState("count", "key2"), (20,))
+
+ def test_delete_list_state(self):
+ """Test that deleteState correctly deletes list state for a given key."""
+ processor = TopKProcessorFactory(k=3).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
+ tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])
+
+ tester.deleteState("topK", "key1")
+ self.assertEqual(tester.peekListState("topK", "key1"), [])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])
+
+ def test_delete_map_state(self):
+ """Test that deleteState correctly deletes map state for a given key."""
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (5,), ("world",): (3,)},
+ )
+
+ tester.deleteState("frequencies", "user1")
+ self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})
+
+ def test_delete_nonexistent_state_raises_error(self):
+ """Test that deleteState raises an error for non-existent state."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
+ tester.deleteState("nonexistent_state", "key1")
+
+ # Timer tests.
+
+ def test_processing_time_timers(self):
+ """Test that TwsTester supports ProcessingTime timers."""
+ processor = SessionTimeoutProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.test("key1", [Row(value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello")])
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(expired1, [])
+
+ # Process input for key2 at t=5000 - should register timer at t=15000.
+ result2 = tester.test("key2", [Row(value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world")])
+
+ # Set processing time to 11000 - key1's timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ self.assertEqual(expired2, [Row(key="key1", result="session-expired")])
+
+ # Set processing time to 16000 - key2's timer should fire.
+ expired3 = tester.setProcessingTime(16000)
+ self.assertEqual(expired3, [Row(key="key2", result="session-expired")])
+
+ # Verify state is cleared after session expiry.
+ self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
+ self.assertIsNone(tester.peekValueState("lastSeen", "key2"))
+
+ def test_processing_time_timers_pandas(self):
+ """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
+ processor = SessionTimeoutProcessorFactory().pandas()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
+ pdt.assert_frame_equal(result1, expected1, check_like=True)
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(len(expired1), 0)
+
+ # Set processing time to 11000 - timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ expected2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
+ pdt.assert_frame_equal(expired2, expected2, check_like=True)
+
+ def test_event_time_timers_manual_watermark(self):
+ """Test that TwsTester supports EventTime timers fired by manual watermark advance."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process event at t=10000 for key1 - registers timer at t=15000.
+ result1 = tester.test("key1", [Row(event_time_ms=10000, value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello@10000")])
+
+ # Process event at t=11000 for key2 - registers timer at t=16000.
+ result2 = tester.test("key2", [Row(event_time_ms=11000, value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world@11000")])
+
+ # Set watermark to 15000 - key1's timer should fire.
+ expired1 = tester.setWatermark(15000)
+ self.assertEqual(len(expired1), 1)
+ self.assertEqual(expired1[0], Row(key="key1", result="session-expired@watermark=15000"))
+
+ # Set watermark to 16000 - key2's timer should fire.
+ expired2 = tester.setWatermark(16000)
+ self.assertEqual(len(expired2), 1)
+ self.assertEqual(expired2[0], Row(key="key2", result="session-expired@watermark=16000"))
+
+ # Verify state is cleared.
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key1"))
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key2"))
+
+ def test_late_event_filtering(self):
+ """Test that TwsTester filters late events in EventTime mode."""
+ processor = EventTimeCountProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Initially watermark is 0, so all events should be processed.
+ result1 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=5000, value="a"),
+ Row(event_time_ms=6000, value="b"),
+ Row(event_time_ms=7000, value="c"),
+ ],
+ )
+ self.assertEqual(result1, [Row(key="key1", count=3)])
+ self.assertEqual(tester.peekValueState("count", "key1"), (3,))
+
+ # Set watermark to 6000.
+ tester.setWatermark(6000)
+
+ # Now watermark is 6000. Send events with mixed event times:
+ # - (4000, "late1") -> event time 4000 <= watermark 6000, should be FILTERED.
+ # - (5000, "late2") -> event time 5000 <= watermark 6000, should be FILTERED.
+ # - (6000, "late3") -> event time 6000 <= watermark 6000, should be FILTERED.
+ # - (8000, "ontime") -> event time 8000 > watermark 6000, should be PROCESSED.
+ result2 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=4000, value="late1"),
+ Row(event_time_ms=5000, value="late2"),
+ Row(event_time_ms=6000, value="late3"),
+ Row(event_time_ms=8000, value="ontime"),
+ ],
+ )
+ # Only 1 event should be processed (the on-time one)
+ self.assertEqual(result2, [Row(key="key1", count=4)]) # 3 + 1 = 4
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+
+ def test_all_late_events_filtered(self):
Review Comment:
nit: I'm not sure this worths a separate test?
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ initial_state_call_count = [0]
+
+ class InitialStateCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ initial_state_call_count[0] += 1
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitialStateCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(initial_state_call_count[0], 0)
+
+ def test_delete_value_state(self):
+ """Test that deleteState correctly deletes value state for a given key."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateValueState("count", "key1", (10,))
+ tester.updateValueState("count", "key2", (20,))
+ self.assertEqual(tester.peekValueState("count", "key1"), (10,))
+
+ tester.deleteState("count", "key1")
+ self.assertIsNone(tester.peekValueState("count", "key1"))
+ self.assertEqual(tester.peekValueState("count", "key2"), (20,))
+
+ def test_delete_list_state(self):
+ """Test that deleteState correctly deletes list state for a given key."""
+ processor = TopKProcessorFactory(k=3).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
+ tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])
+
+ tester.deleteState("topK", "key1")
+ self.assertEqual(tester.peekListState("topK", "key1"), [])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])
+
+ def test_delete_map_state(self):
+ """Test that deleteState correctly deletes map state for a given key."""
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (5,), ("world",): (3,)},
+ )
+
+ tester.deleteState("frequencies", "user1")
+ self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})
+
+ def test_delete_nonexistent_state_raises_error(self):
+ """Test that deleteState raises an error for non-existent state."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
+ tester.deleteState("nonexistent_state", "key1")
+
+ # Timer tests.
+
+ def test_processing_time_timers(self):
+ """Test that TwsTester supports ProcessingTime timers."""
+ processor = SessionTimeoutProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.test("key1", [Row(value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello")])
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(expired1, [])
+
+ # Process input for key2 at t=5000 - should register timer at t=15000.
+ result2 = tester.test("key2", [Row(value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world")])
+
+ # Set processing time to 11000 - key1's timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ self.assertEqual(expired2, [Row(key="key1", result="session-expired")])
+
+ # Set processing time to 16000 - key2's timer should fire.
+ expired3 = tester.setProcessingTime(16000)
+ self.assertEqual(expired3, [Row(key="key2", result="session-expired")])
+
+ # Verify state is cleared after session expiry.
+ self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
+ self.assertIsNone(tester.peekValueState("lastSeen", "key2"))
+
+ def test_processing_time_timers_pandas(self):
+ """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
+ processor = SessionTimeoutProcessorFactory().pandas()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
+ pdt.assert_frame_equal(result1, expected1, check_like=True)
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(len(expired1), 0)
+
+ # Set processing time to 11000 - timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ expected2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
+ pdt.assert_frame_equal(expired2, expected2, check_like=True)
+
+ def test_event_time_timers_manual_watermark(self):
+ """Test that TwsTester supports EventTime timers fired by manual watermark advance."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process event at t=10000 for key1 - registers timer at t=15000.
+ result1 = tester.test("key1", [Row(event_time_ms=10000, value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello@10000")])
+
+ # Process event at t=11000 for key2 - registers timer at t=16000.
+ result2 = tester.test("key2", [Row(event_time_ms=11000, value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world@11000")])
+
+ # Set watermark to 15000 - key1's timer should fire.
+ expired1 = tester.setWatermark(15000)
+ self.assertEqual(len(expired1), 1)
+ self.assertEqual(expired1[0], Row(key="key1", result="session-expired@watermark=15000"))
+
+ # Set watermark to 16000 - key2's timer should fire.
+ expired2 = tester.setWatermark(16000)
+ self.assertEqual(len(expired2), 1)
+ self.assertEqual(expired2[0], Row(key="key2", result="session-expired@watermark=16000"))
+
+ # Verify state is cleared.
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key1"))
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key2"))
+
+ def test_late_event_filtering(self):
Review Comment:
nit: It'd probably easier to review if we just keep the same logic among pair of tests. The logic is simple enough so wasn't a big deal but was doing side-by-side comparison likewise other test cases and realized it's different. No need to change to be in sync though.
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
Review Comment:
nit: Not very sure this is needed, but this might be arguable so OK to leave it as it is.
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ initial_state_call_count = [0]
+
+ class InitialStateCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ initial_state_call_count[0] += 1
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitialStateCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(initial_state_call_count[0], 0)
+
+ def test_delete_value_state(self):
+ """Test that deleteState correctly deletes value state for a given key."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateValueState("count", "key1", (10,))
+ tester.updateValueState("count", "key2", (20,))
+ self.assertEqual(tester.peekValueState("count", "key1"), (10,))
+
+ tester.deleteState("count", "key1")
+ self.assertIsNone(tester.peekValueState("count", "key1"))
+ self.assertEqual(tester.peekValueState("count", "key2"), (20,))
+
+ def test_delete_list_state(self):
+ """Test that deleteState correctly deletes list state for a given key."""
+ processor = TopKProcessorFactory(k=3).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
+ tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])
+
+ tester.deleteState("topK", "key1")
+ self.assertEqual(tester.peekListState("topK", "key1"), [])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])
+
+ def test_delete_map_state(self):
+ """Test that deleteState correctly deletes map state for a given key."""
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (5,), ("world",): (3,)},
+ )
+
+ tester.deleteState("frequencies", "user1")
+ self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})
+
+ def test_delete_nonexistent_state_raises_error(self):
+ """Test that deleteState raises an error for non-existent state."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
+ tester.deleteState("nonexistent_state", "key1")
+
+ # Timer tests.
+
+ def test_processing_time_timers(self):
+ """Test that TwsTester supports ProcessingTime timers."""
+ processor = SessionTimeoutProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.test("key1", [Row(value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello")])
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(expired1, [])
+
+ # Process input for key2 at t=5000 - should register timer at t=15000.
+ result2 = tester.test("key2", [Row(value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world")])
+
+ # Set processing time to 11000 - key1's timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ self.assertEqual(expired2, [Row(key="key1", result="session-expired")])
+
+ # Set processing time to 16000 - key2's timer should fire.
+ expired3 = tester.setProcessingTime(16000)
+ self.assertEqual(expired3, [Row(key="key2", result="session-expired")])
+
+ # Verify state is cleared after session expiry.
+ self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
+ self.assertIsNone(tester.peekValueState("lastSeen", "key2"))
+
+ def test_processing_time_timers_pandas(self):
+ """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
+ processor = SessionTimeoutProcessorFactory().pandas()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
+ pdt.assert_frame_equal(result1, expected1, check_like=True)
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(len(expired1), 0)
+
+ # Set processing time to 11000 - timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ expected2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
+ pdt.assert_frame_equal(expired2, expected2, check_like=True)
+
+ def test_event_time_timers_manual_watermark(self):
Review Comment:
nit: most tests are doing the same thing with pair of Scala test but this doesn't seem to do the same - cancellation of timer isn't tested here. Is it missing?
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ initial_state_call_count = [0]
+
+ class InitialStateCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ initial_state_call_count[0] += 1
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitialStateCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(initial_state_call_count[0], 0)
+
+ def test_delete_value_state(self):
+ """Test that deleteState correctly deletes value state for a given key."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateValueState("count", "key1", (10,))
+ tester.updateValueState("count", "key2", (20,))
+ self.assertEqual(tester.peekValueState("count", "key1"), (10,))
+
+ tester.deleteState("count", "key1")
+ self.assertIsNone(tester.peekValueState("count", "key1"))
+ self.assertEqual(tester.peekValueState("count", "key2"), (20,))
+
+ def test_delete_list_state(self):
+ """Test that deleteState correctly deletes list state for a given key."""
+ processor = TopKProcessorFactory(k=3).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
+ tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])
+
+ tester.deleteState("topK", "key1")
+ self.assertEqual(tester.peekListState("topK", "key1"), [])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])
+
+ def test_delete_map_state(self):
+ """Test that deleteState correctly deletes map state for a given key."""
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (5,), ("world",): (3,)},
+ )
+
+ tester.deleteState("frequencies", "user1")
+ self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})
+
+ def test_delete_nonexistent_state_raises_error(self):
+ """Test that deleteState raises an error for non-existent state."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
+ tester.deleteState("nonexistent_state", "key1")
+
+ # Timer tests.
+
+ def test_processing_time_timers(self):
+ """Test that TwsTester supports ProcessingTime timers."""
+ processor = SessionTimeoutProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.test("key1", [Row(value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello")])
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(expired1, [])
+
+ # Process input for key2 at t=5000 - should register timer at t=15000.
+ result2 = tester.test("key2", [Row(value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world")])
+
+ # Set processing time to 11000 - key1's timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ self.assertEqual(expired2, [Row(key="key1", result="session-expired")])
+
+ # Set processing time to 16000 - key2's timer should fire.
+ expired3 = tester.setProcessingTime(16000)
+ self.assertEqual(expired3, [Row(key="key2", result="session-expired")])
+
+ # Verify state is cleared after session expiry.
+ self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
+ self.assertIsNone(tester.peekValueState("lastSeen", "key2"))
+
+ def test_processing_time_timers_pandas(self):
+ """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
+ processor = SessionTimeoutProcessorFactory().pandas()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
+ pdt.assert_frame_equal(result1, expected1, check_like=True)
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(len(expired1), 0)
+
+ # Set processing time to 11000 - timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ expected2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
+ pdt.assert_frame_equal(expired2, expected2, check_like=True)
+
+ def test_event_time_timers_manual_watermark(self):
+ """Test that TwsTester supports EventTime timers fired by manual watermark advance."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process event at t=10000 for key1 - registers timer at t=15000.
+ result1 = tester.test("key1", [Row(event_time_ms=10000, value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello@10000")])
+
+ # Process event at t=11000 for key2 - registers timer at t=16000.
+ result2 = tester.test("key2", [Row(event_time_ms=11000, value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world@11000")])
+
+ # Set watermark to 15000 - key1's timer should fire.
+ expired1 = tester.setWatermark(15000)
+ self.assertEqual(len(expired1), 1)
+ self.assertEqual(expired1[0], Row(key="key1", result="session-expired@watermark=15000"))
+
+ # Set watermark to 16000 - key2's timer should fire.
+ expired2 = tester.setWatermark(16000)
+ self.assertEqual(len(expired2), 1)
+ self.assertEqual(expired2[0], Row(key="key2", result="session-expired@watermark=16000"))
+
+ # Verify state is cleared.
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key1"))
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key2"))
+
+ def test_late_event_filtering(self):
+ """Test that TwsTester filters late events in EventTime mode."""
+ processor = EventTimeCountProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Initially watermark is 0, so all events should be processed.
+ result1 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=5000, value="a"),
+ Row(event_time_ms=6000, value="b"),
+ Row(event_time_ms=7000, value="c"),
+ ],
+ )
+ self.assertEqual(result1, [Row(key="key1", count=3)])
+ self.assertEqual(tester.peekValueState("count", "key1"), (3,))
+
+ # Set watermark to 6000.
+ tester.setWatermark(6000)
+
+ # Now watermark is 6000. Send events with mixed event times:
+ # - (4000, "late1") -> event time 4000 <= watermark 6000, should be FILTERED.
+ # - (5000, "late2") -> event time 5000 <= watermark 6000, should be FILTERED.
+ # - (6000, "late3") -> event time 6000 <= watermark 6000, should be FILTERED.
+ # - (8000, "ontime") -> event time 8000 > watermark 6000, should be PROCESSED.
+ result2 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=4000, value="late1"),
+ Row(event_time_ms=5000, value="late2"),
+ Row(event_time_ms=6000, value="late3"),
+ Row(event_time_ms=8000, value="ontime"),
+ ],
+ )
+ # Only 1 event should be processed (the on-time one)
+ self.assertEqual(result2, [Row(key="key1", count=4)]) # 3 + 1 = 4
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+
+ def test_all_late_events_filtered(self):
+ """Test that all late events are filtered when all are older than watermark."""
+ processor = EventTimeCountProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process events to get initial count.
+ tester.test("key1", [Row(event_time_ms=10000, value="a")])
+ self.assertEqual(tester.peekValueState("count", "key1"), (1,))
+
+ # Set watermark to 10000.
+ tester.setWatermark(10000)
+
+ # Now send only late events - all should be filtered, count unchanged.
+ result = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=1000, value="late1"),
+ Row(event_time_ms=5000, value="late2"),
+ Row(event_time_ms=9999, value="late3"),
+ Row(event_time_ms=10000, value="late4"), # exactly at watermark, also filtered
+ ],
+ )
+ # No events processed - handleInputRows is called with empty list.
+ self.assertEqual(result, [Row(key="key1", count=1)]) # Count still 1
+ self.assertEqual(tester.peekValueState("count", "key1"), (1,))
+
+ def test_timer_registration_in_none_mode_raises_error(self):
+ """Test that timer operations raise error in TimeMode.None."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_set_processing_time_in_wrong_mode_raises_error(self):
+ """Test that setProcessingTime raises error in non-ProcessingTime mode."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ # Test in None mode.
+ tester_none = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_none.setProcessingTime(1000)
+
+ def test_set_watermark_in_wrong_mode_raises_error(self):
+ """Test that setWatermark raises error in non-EventTime mode."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ # Test in None mode.
+ tester_none = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_none.setWatermark(1000)
+
+ # Test in ProcessingTime mode.
+ tester_pt = TwsTester(processor, timeMode="ProcessingTime")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_pt.setWatermark(1000)
+
+ def test_eventtime_mode_requires_extractor(self):
+ """Test that EventTime mode requires eventTimeExtractor."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ with self.assertRaisesRegex(PySparkValueError, "ARGUMENT_REQUIRED"):
+ TwsTester(processor, timeMode="EventTime")
+
+ def test_processing_time_must_move_forward(self):
+ """Test that setProcessingTime requires time to move forward."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ tester.setProcessingTime(1000)
+
+ # Same time should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setProcessingTime(1000)
+
+ # Earlier time should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setProcessingTime(500)
+
+ def test_watermark_must_move_forward(self):
+ """Test that setWatermark requires watermark to move forward."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ tester.setWatermark(1000)
+
+ # Same watermark should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setWatermark(1000)
+
+ # Earlier watermark should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setWatermark(500)
+
+ def test_timer_cannot_be_registered_in_past(self):
+ """Test that timer cannot be registered at or before current watermark in EventTime mode."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Set watermark to 10000.
+ tester.setWatermark(10000)
+
+ # Try to register a timer at or before watermark - should raise error.
+ tester.handle.setGroupingKey("key1")
+ with self.assertRaises(IllegalArgumentException):
+ tester.handle.registerTimer(10000) # exactly at watermark
+
+ with self.assertRaises(IllegalArgumentException):
+ tester.handle.registerTimer(5000) # before watermark
+
Review Comment:
Do we have pair test here for "TwsTester with eventTimeExtractor should allow timer registration from on-time events"? If not is there a reason or is it just missing?
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
Review Comment:
nit: I don't think this test is useful.
Unfortunately this type of issues has happened multiple times in this test suite. I'd be the first code reviewer of the code change LLM generates if I submit a PR - the code authorship is given to PR author, not LLM. Code quality is not guaranteed by LLM - human should take responsibility.
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ initial_state_call_count = [0]
+
+ class InitialStateCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ initial_state_call_count[0] += 1
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitialStateCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(initial_state_call_count[0], 0)
+
+ def test_delete_value_state(self):
+ """Test that deleteState correctly deletes value state for a given key."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateValueState("count", "key1", (10,))
+ tester.updateValueState("count", "key2", (20,))
+ self.assertEqual(tester.peekValueState("count", "key1"), (10,))
+
+ tester.deleteState("count", "key1")
+ self.assertIsNone(tester.peekValueState("count", "key1"))
+ self.assertEqual(tester.peekValueState("count", "key2"), (20,))
+
+ def test_delete_list_state(self):
+ """Test that deleteState correctly deletes list state for a given key."""
+ processor = TopKProcessorFactory(k=3).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
+ tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])
+
+ tester.deleteState("topK", "key1")
+ self.assertEqual(tester.peekListState("topK", "key1"), [])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])
+
+ def test_delete_map_state(self):
+ """Test that deleteState correctly deletes map state for a given key."""
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (5,), ("world",): (3,)},
+ )
+
+ tester.deleteState("frequencies", "user1")
+ self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})
+
+ def test_delete_nonexistent_state_raises_error(self):
+ """Test that deleteState raises an error for non-existent state."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
+ tester.deleteState("nonexistent_state", "key1")
+
+ # Timer tests.
+
+ def test_processing_time_timers(self):
+ """Test that TwsTester supports ProcessingTime timers."""
+ processor = SessionTimeoutProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.test("key1", [Row(value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello")])
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(expired1, [])
+
+ # Process input for key2 at t=5000 - should register timer at t=15000.
+ result2 = tester.test("key2", [Row(value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world")])
+
+ # Set processing time to 11000 - key1's timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ self.assertEqual(expired2, [Row(key="key1", result="session-expired")])
+
+ # Set processing time to 16000 - key2's timer should fire.
+ expired3 = tester.setProcessingTime(16000)
+ self.assertEqual(expired3, [Row(key="key2", result="session-expired")])
+
+ # Verify state is cleared after session expiry.
+ self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
+ self.assertIsNone(tester.peekValueState("lastSeen", "key2"))
+
+ def test_processing_time_timers_pandas(self):
+ """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
+ processor = SessionTimeoutProcessorFactory().pandas()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
+ pdt.assert_frame_equal(result1, expected1, check_like=True)
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(len(expired1), 0)
+
+ # Set processing time to 11000 - timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ expected2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
+ pdt.assert_frame_equal(expired2, expected2, check_like=True)
+
+ def test_event_time_timers_manual_watermark(self):
+ """Test that TwsTester supports EventTime timers fired by manual watermark advance."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process event at t=10000 for key1 - registers timer at t=15000.
+ result1 = tester.test("key1", [Row(event_time_ms=10000, value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello@10000")])
+
+ # Process event at t=11000 for key2 - registers timer at t=16000.
+ result2 = tester.test("key2", [Row(event_time_ms=11000, value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world@11000")])
+
+ # Set watermark to 15000 - key1's timer should fire.
+ expired1 = tester.setWatermark(15000)
+ self.assertEqual(len(expired1), 1)
+ self.assertEqual(expired1[0], Row(key="key1", result="session-expired@watermark=15000"))
+
+ # Set watermark to 16000 - key2's timer should fire.
+ expired2 = tester.setWatermark(16000)
+ self.assertEqual(len(expired2), 1)
+ self.assertEqual(expired2[0], Row(key="key2", result="session-expired@watermark=16000"))
+
+ # Verify state is cleared.
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key1"))
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key2"))
+
+ def test_late_event_filtering(self):
+ """Test that TwsTester filters late events in EventTime mode."""
+ processor = EventTimeCountProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Initially watermark is 0, so all events should be processed.
+ result1 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=5000, value="a"),
+ Row(event_time_ms=6000, value="b"),
+ Row(event_time_ms=7000, value="c"),
+ ],
+ )
+ self.assertEqual(result1, [Row(key="key1", count=3)])
+ self.assertEqual(tester.peekValueState("count", "key1"), (3,))
+
+ # Set watermark to 6000.
+ tester.setWatermark(6000)
+
+ # Now watermark is 6000. Send events with mixed event times:
+ # - (4000, "late1") -> event time 4000 <= watermark 6000, should be FILTERED.
+ # - (5000, "late2") -> event time 5000 <= watermark 6000, should be FILTERED.
+ # - (6000, "late3") -> event time 6000 <= watermark 6000, should be FILTERED.
+ # - (8000, "ontime") -> event time 8000 > watermark 6000, should be PROCESSED.
+ result2 = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=4000, value="late1"),
+ Row(event_time_ms=5000, value="late2"),
+ Row(event_time_ms=6000, value="late3"),
+ Row(event_time_ms=8000, value="ontime"),
+ ],
+ )
+ # Only 1 event should be processed (the on-time one)
+ self.assertEqual(result2, [Row(key="key1", count=4)]) # 3 + 1 = 4
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+
+ def test_all_late_events_filtered(self):
+ """Test that all late events are filtered when all are older than watermark."""
+ processor = EventTimeCountProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process events to get initial count.
+ tester.test("key1", [Row(event_time_ms=10000, value="a")])
+ self.assertEqual(tester.peekValueState("count", "key1"), (1,))
+
+ # Set watermark to 10000.
+ tester.setWatermark(10000)
+
+ # Now send only late events - all should be filtered, count unchanged.
+ result = tester.test(
+ "key1",
+ [
+ Row(event_time_ms=1000, value="late1"),
+ Row(event_time_ms=5000, value="late2"),
+ Row(event_time_ms=9999, value="late3"),
+ Row(event_time_ms=10000, value="late4"), # exactly at watermark, also filtered
+ ],
+ )
+ # No events processed - handleInputRows is called with empty list.
+ self.assertEqual(result, [Row(key="key1", count=1)]) # Count still 1
+ self.assertEqual(tester.peekValueState("count", "key1"), (1,))
+
+ def test_timer_registration_in_none_mode_raises_error(self):
+ """Test that timer operations raise error in TimeMode.None."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_set_processing_time_in_wrong_mode_raises_error(self):
+ """Test that setProcessingTime raises error in non-ProcessingTime mode."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ # Test in None mode.
+ tester_none = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_none.setProcessingTime(1000)
+
+ def test_set_watermark_in_wrong_mode_raises_error(self):
+ """Test that setWatermark raises error in non-EventTime mode."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ # Test in None mode.
+ tester_none = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_none.setWatermark(1000)
+
+ # Test in ProcessingTime mode.
+ tester_pt = TwsTester(processor, timeMode="ProcessingTime")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester_pt.setWatermark(1000)
+
+ def test_eventtime_mode_requires_extractor(self):
+ """Test that EventTime mode requires eventTimeExtractor."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ with self.assertRaisesRegex(PySparkValueError, "ARGUMENT_REQUIRED"):
+ TwsTester(processor, timeMode="EventTime")
+
+ def test_processing_time_must_move_forward(self):
+ """Test that setProcessingTime requires time to move forward."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ tester.setProcessingTime(1000)
+
+ # Same time should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setProcessingTime(1000)
+
+ # Earlier time should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setProcessingTime(500)
+
+ def test_watermark_must_move_forward(self):
+ """Test that setWatermark requires watermark to move forward."""
+ processor = RunningCountStatefulProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ tester.setWatermark(1000)
+
+ # Same watermark should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setWatermark(1000)
+
+ # Earlier watermark should raise error.
+ with self.assertRaises(IllegalArgumentException):
+ tester.setWatermark(500)
+
+ def test_timer_cannot_be_registered_in_past(self):
+ """Test that timer cannot be registered at or before current watermark in EventTime mode."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Set watermark to 10000.
+ tester.setWatermark(10000)
+
+ # Try to register a timer at or before watermark - should raise error.
+ tester.handle.setGroupingKey("key1")
+ with self.assertRaises(IllegalArgumentException):
+ tester.handle.registerTimer(10000) # exactly at watermark
+
+ with self.assertRaises(IllegalArgumentException):
+ tester.handle.registerTimer(5000) # before watermark
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterFuzzTests(ReusedSQLTestCase):
+ @classmethod
+ def conf(cls):
+ cfg = SparkConf()
+ cfg.set("spark.sql.shuffle.partitions", "1")
+ cfg.set(
+ "spark.sql.streaming.stateStore.providerClass",
+ "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider",
+ )
+ return cfg
+
+ # Helper that runs real streaming query with TWS and compares results with TwsTester output.
+ # Supports both row mode (use_pandas=False) and Pandas mode (use_pandas=True).
+ # Supports multiple batches.
+ # Assumes that the first column is the key column.
+ # Output rows are checked ignoring order.
+ # Correct only if results do not depend on reordering rows within batch.
+ def _check_tws_tester(
+ self,
+ processor: StatefulProcessor,
+ batches: list[list[tuple]],
+ input_columns: list[str],
+ input_types: list[str],
+ output_schema: StructType,
+ use_pandas: bool = False,
+ ):
+ num_batches = len(batches)
+ key_column: str = input_columns[0] # Assume first column is the key column.
+ value_columns: list[str] = input_columns[1:] # All columns except the key column.
+
+ num_input_columns: int = len(input_columns)
+ assert len(input_types) == num_input_columns
+
+ # Use TwsTester to compute expected results.
+ tester = TwsTester(processor)
+
+ def _run_tester(batch: list[tuple]) -> list[Row]:
+ # Group batch by key.
+ grouped = {}
+ for row in batch:
+ key = row[0]
+ if key not in grouped:
+ grouped[key] = []
+ grouped[key].append(row[1:])
+
+ results: list[Row] = []
+ for key, values in grouped.items():
+ if use_pandas:
+ input_to_tester = pd.DataFrame(values, columns=value_columns)
+ result_df = tester.testInPandas(key, input_to_tester)
+ results.extend(self.spark.createDataFrame(result_df).collect())
+ else:
+ input_to_tester: list[Row] = [
+ Row(**dict(zip(value_columns, v))) for v in values
+ ]
+ results.extend(tester.test(key, input_to_tester))
+ return results
+
+ expected_results: list[list[Row]] = [_run_tester(batch) for batch in batches]
+
+ # Create streaming DataFrame, it will read input from files.
+ input_path = tempfile.mkdtemp()
+ df = self.spark.readStream.format("text").load(input_path)
+ df_split = df.withColumn("split_values", split(df["value"], ","))
+ df_final = df_split.select(
+ *[
+ df_split.split_values.getItem(i).alias(col_name).cast(input_types[i])
+ for i, col_name in enumerate(input_columns)
+ ]
+ )
+ self.assertTrue(df_final.isStreaming)
+
+ # Apply TransformWithState.
+ df_grouped = df_final.groupBy(key_column)
+ if use_pandas:
+ tws_df = df_grouped.transformWithStateInPandas(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ )
+ else:
+ tws_df = df_grouped.transformWithState(
+ statefulProcessor=processor,
+ outputStructType=output_schema,
+ outputMode="Update",
+ timeMode="None",
+ )
+
+ # Start streaming query collecting TWS results.
+ actual_results: list[list[Row]] = []
+
+ def _collect_results(batch_df: DataFrame, batch_id):
+ rows: list[Row] = batch_df.collect()
+ actual_results.append(rows)
+
+ query: StreamingQuery = (
+ tws_df.writeStream.queryName("test_tmp_query")
+ .foreachBatch(_collect_results)
+ .outputMode("update")
+ .start()
+ )
+
+ # Process each input batch.
+ def _write_batch(batch_id: int, batch: list[tuple]):
+ with open(input_path + f"/test{batch_id}.txt", "w") as f:
+ for row in batch:
+ assert len(row) == num_input_columns
+ f.write(",".join(str(value) for value in row) + "\n")
+
+ try:
+ for batch_id, batch in enumerate(batches):
+ _write_batch(batch_id, batch)
+ query.processAllAvailable()
+ query.awaitTermination(1)
+ self.assertTrue(query.exception() is None)
+ finally:
+ query.stop()
+
+ # Assert actual results are equal to expected results.
+ self.assertEqual(len(actual_results), num_batches)
+ self.assertEqual(len(expected_results), num_batches)
+ for batch_id in range(num_batches):
+ self.assertCountEqual(actual_results[batch_id], expected_results[batch_id])
Review Comment:
(Not a comment about the code change...)
Wow, the method name `assertCountEqual` is actually messed up... This gives me impression we only check the size and I was about to raise concern, and I realized it's just a comparison without ordering.
--
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #53758:
URL: https://github.com/apache/spark/pull/53758#issuecomment-3850959430
https://github.com/fedimser/spark/actions/runs/21688225187/job/62553376983
https://github.com/fedimser/spark/actions/runs/21692971251/job/62556992226
It seems like the failure of modules are different which backs up the claim. Since we gave two times of trials, I'd love to avoid asking to trigger more.
--
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r2756166663
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ initial_state_call_count = [0]
+
+ class InitialStateCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ initial_state_call_count[0] += 1
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitialStateCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(initial_state_call_count[0], 0)
+
+ def test_delete_value_state(self):
+ """Test that deleteState correctly deletes value state for a given key."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateValueState("count", "key1", (10,))
+ tester.updateValueState("count", "key2", (20,))
+ self.assertEqual(tester.peekValueState("count", "key1"), (10,))
+
+ tester.deleteState("count", "key1")
+ self.assertIsNone(tester.peekValueState("count", "key1"))
+ self.assertEqual(tester.peekValueState("count", "key2"), (20,))
+
+ def test_delete_list_state(self):
+ """Test that deleteState correctly deletes list state for a given key."""
+ processor = TopKProcessorFactory(k=3).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
+ tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])
+
+ tester.deleteState("topK", "key1")
+ self.assertEqual(tester.peekListState("topK", "key1"), [])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])
+
+ def test_delete_map_state(self):
+ """Test that deleteState correctly deletes map state for a given key."""
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (5,), ("world",): (3,)},
+ )
+
+ tester.deleteState("frequencies", "user1")
+ self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})
+
+ def test_delete_nonexistent_state_raises_error(self):
+ """Test that deleteState raises an error for non-existent state."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
+ tester.deleteState("nonexistent_state", "key1")
+
+ # Timer tests.
+
+ def test_processing_time_timers(self):
+ """Test that TwsTester supports ProcessingTime timers."""
+ processor = SessionTimeoutProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.test("key1", [Row(value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello")])
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(expired1, [])
+
+ # Process input for key2 at t=5000 - should register timer at t=15000.
+ result2 = tester.test("key2", [Row(value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world")])
+
+ # Set processing time to 11000 - key1's timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ self.assertEqual(expired2, [Row(key="key1", result="session-expired")])
+
+ # Set processing time to 16000 - key2's timer should fire.
+ expired3 = tester.setProcessingTime(16000)
+ self.assertEqual(expired3, [Row(key="key2", result="session-expired")])
+
+ # Verify state is cleared after session expiry.
+ self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
+ self.assertIsNone(tester.peekValueState("lastSeen", "key2"))
+
+ def test_processing_time_timers_pandas(self):
+ """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
+ processor = SessionTimeoutProcessorFactory().pandas()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
+ pdt.assert_frame_equal(result1, expected1, check_like=True)
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(len(expired1), 0)
+
+ # Set processing time to 11000 - timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ expected2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
+ pdt.assert_frame_equal(expired2, expected2, check_like=True)
+
+ def test_event_time_timers_manual_watermark(self):
+ """Test that TwsTester supports EventTime timers fired by manual watermark advance."""
+ processor = EventTimeSessionProcessorFactory().row()
+
+ def event_time_extractor(row: Row) -> int:
+ return row.event_time_ms
+
+ tester = TwsTester(
+ processor,
+ timeMode="EventTime",
+ eventTimeExtractor=event_time_extractor,
+ )
+
+ # Process event at t=10000 for key1 - registers timer at t=15000.
+ result1 = tester.test("key1", [Row(event_time_ms=10000, value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello@10000")])
+
+ # Process event at t=11000 for key2 - registers timer at t=16000.
+ result2 = tester.test("key2", [Row(event_time_ms=11000, value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world@11000")])
+
+ # Set watermark to 15000 - key1's timer should fire.
+ expired1 = tester.setWatermark(15000)
+ self.assertEqual(len(expired1), 1)
+ self.assertEqual(expired1[0], Row(key="key1", result="session-expired@watermark=15000"))
+
+ # Set watermark to 16000 - key2's timer should fire.
+ expired2 = tester.setWatermark(16000)
+ self.assertEqual(len(expired2), 1)
+ self.assertEqual(expired2[0], Row(key="key2", result="session-expired@watermark=16000"))
+
+ # Verify state is cleared.
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key1"))
+ self.assertIsNone(tester.peekValueState("lastEventTime", "key2"))
+
+ def test_late_event_filtering(self):
Review Comment:
It is mostly in sync now.
There are some differences that you pointed out, e.g. here we cover null keys or clearing nonexistent state. I will not remove these tests just because we don't have them in Scala version.
Also some tests are duplicated to test Row/Pandas API, but I didn't do this for all tests, the decision which tests to duplicate was arbitrary, I just wanted to cover all code paths.
--
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r2755792522
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
Review Comment:
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r2756104956
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ result = tester.test("key1", [Row(value="b")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_handle_initial_state_not_called_without_initial_state(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ initial_state_call_count = [0]
+
+ class InitialStateCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInitialState(self, key, initialState, timerValues) -> None:
+ initial_state_call_count[0] += 1
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitialStateCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(initial_state_call_count[0], 0)
+
+ def test_delete_value_state(self):
+ """Test that deleteState correctly deletes value state for a given key."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateValueState("count", "key1", (10,))
+ tester.updateValueState("count", "key2", (20,))
+ self.assertEqual(tester.peekValueState("count", "key1"), (10,))
+
+ tester.deleteState("count", "key1")
+ self.assertIsNone(tester.peekValueState("count", "key1"))
+ self.assertEqual(tester.peekValueState("count", "key2"), (20,))
+
+ def test_delete_list_state(self):
+ """Test that deleteState correctly deletes list state for a given key."""
+ processor = TopKProcessorFactory(k=3).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "key1", [(1.0,), (2.0,), (3.0,)])
+ tester.updateListState("topK", "key2", [(4.0,), (5.0,)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(1.0,), (2.0,), (3.0,)])
+
+ tester.deleteState("topK", "key1")
+ self.assertEqual(tester.peekListState("topK", "key1"), [])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(4.0,), (5.0,)])
+
+ def test_delete_map_state(self):
+ """Test that deleteState correctly deletes map state for a given key."""
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (5,), ("world",): (3,)},
+ )
+
+ tester.deleteState("frequencies", "user1")
+ self.assertEqual(tester.peekMapState("frequencies", "user1"), {})
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (10,)})
+
+ def test_delete_nonexistent_state_raises_error(self):
+ """Test that deleteState raises an error for non-existent state."""
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ with self.assertRaisesRegex(PySparkAssertionError, "STATE_NOT_EXISTS"):
+ tester.deleteState("nonexistent_state", "key1")
+
+ # Timer tests.
+
+ def test_processing_time_timers(self):
+ """Test that TwsTester supports ProcessingTime timers."""
+ processor = SessionTimeoutProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.test("key1", [Row(value="hello")])
+ self.assertEqual(result1, [Row(key="key1", result="received:hello")])
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(expired1, [])
+
+ # Process input for key2 at t=5000 - should register timer at t=15000.
+ result2 = tester.test("key2", [Row(value="world")])
+ self.assertEqual(result2, [Row(key="key2", result="received:world")])
+
+ # Set processing time to 11000 - key1's timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ self.assertEqual(expired2, [Row(key="key1", result="session-expired")])
+
+ # Set processing time to 16000 - key2's timer should fire.
+ expired3 = tester.setProcessingTime(16000)
+ self.assertEqual(expired3, [Row(key="key2", result="session-expired")])
+
+ # Verify state is cleared after session expiry.
+ self.assertIsNone(tester.peekValueState("lastSeen", "key1"))
+ self.assertIsNone(tester.peekValueState("lastSeen", "key2"))
+
+ def test_processing_time_timers_pandas(self):
+ """Test that TwsTester supports ProcessingTime timers in Pandas mode."""
+ processor = SessionTimeoutProcessorFactory().pandas()
+ tester = TwsTester(processor, timeMode="ProcessingTime")
+
+ # Process input for key1 - should register a timer at t=10000.
+ result1 = tester.testInPandas("key1", pd.DataFrame({"value": ["hello"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "result": ["received:hello"]})
+ pdt.assert_frame_equal(result1, expected1, check_like=True)
+
+ # Set processing time to 5000 - timer should NOT fire yet.
+ expired1 = tester.setProcessingTime(5000)
+ self.assertEqual(len(expired1), 0)
+
+ # Set processing time to 11000 - timer should fire.
+ expired2 = tester.setProcessingTime(11000)
+ expected2 = pd.DataFrame({"key": ["key1"], "result": ["session-expired"]})
+ pdt.assert_frame_equal(expired2, expected2, check_like=True)
+
+ def test_event_time_timers_manual_watermark(self):
Review Comment:
There are 2 versions of this test, for Row (above) and Pandas, that both correspond to "TwsTester should support ProcessingTime timers" in Scala.
It seems I wanted to keep the second test shorter.
I extended this test so it does exactly the same thing as the test above, but using the testInPandas API.
--
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #53758:
URL: https://github.com/apache/spark/pull/53758#issuecomment-3849517683
Awesome, thanks for doing that. I roughly remember we tried to make pandas be optional in PySpark code, so if we can do that as well, we should do that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r2755787277
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiValueStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.state1 = handle.getValueState("state1", "int")
+ self.state2 = handle.getValueState("state2", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val1 = self.state1.get() if self.state1.exists() else None
+ count1 = val1[0] if val1 else 0
+ val2 = self.state2.get() if self.state2.exists() else None
+ count2 = val2[0] if val2 else 0
+ for row in rows:
+ count1 += 1
+ count2 += 2
+ self.state1.update((count1,))
+ self.state2.update((count2,))
+ yield Row(key=key[0], count1=count1, count2=count2)
+
+ processor = MultiValueStateProcessor()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count1=1, count2=2)])
+ self.assertEqual(tester.peekValueState("state1", "key1"), (1,))
+ self.assertEqual(tester.peekValueState("state2", "key1"), (2,))
+
+ def test_multiple_list_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiListStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.list1 = handle.getListState("list1", "string")
+ self.list2 = handle.getListState("list2", "string")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ self.list1.appendValue((row.value,))
+ self.list2.appendValue((row.value + "_2",))
+ yield Row(key=key[0], count=1)
+
+ processor = MultiListStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekListState("list1", "key1"), [("a",)])
+ self.assertEqual(tester.peekListState("list2", "key1"), [("a_2",)])
+
+ def test_multiple_map_states(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MultiMapStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.map1 = handle.getMapState("map1", "string", "int")
+ self.map2 = handle.getMapState("map2", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val1 = self.map1.getValue((row.word,))
+ count1 = val1[0] if val1 is not None else 0
+ val2 = self.map2.getValue((row.word,))
+ count2 = val2[0] if val2 is not None else 0
+ self.map1.updateValue((row.word,), (count1 + 1,))
+ self.map2.updateValue((row.word,), (count2 + 2,))
+ yield Row(key=key[0], word=row.word)
+
+ processor = MultiMapStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(word="hello")])
+ self.assertEqual(tester.peekMapState("map1", "key1"), {("hello",): (1,)})
+ self.assertEqual(tester.peekMapState("map2", "key1"), {("hello",): (2,)})
+
+ def test_mixed_state_types(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ class MixedStateProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ self.handle = handle
+ self.value_state = handle.getValueState("value", "int")
+ self.list_state = handle.getListState("list", "string")
+ self.map_state = handle.getMapState("map", "string", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ for row in rows:
+ val = self.value_state.get() if self.value_state.exists() else None
+ count = val[0] if val else 0
+ self.value_state.update((count + 1,))
+ self.list_state.appendValue((row.value,))
+ self.map_state.updateValue((row.value,), (1,))
+ yield Row(key=key[0], count=1)
+
+ processor = MixedStateProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ self.assertEqual(tester.peekValueState("value", "key1"), (1,))
+ self.assertEqual(tester.peekListState("list", "key1"), [("a",)])
+ self.assertEqual(tester.peekMapState("map", "key1"), {("a",): (1,)})
+
+ def test_pandas_with_multiple_rows(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b", "c"]})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [3]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_pandas_dtype_preservation(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a"]})
+ result = tester.testInPandas("key1", input_df)
+ self.assertEqual(result["count"].dtype, "int64")
+
+ def test_processor_init_called_once(self):
+ from pyspark.sql.streaming.stateful_processor import (
+ StatefulProcessor,
+ StatefulProcessorHandle,
+ )
+ from typing import Iterator
+
+ init_call_count = [0]
+
+ class InitCountingProcessor(StatefulProcessor):
+ def init(self, handle: StatefulProcessorHandle) -> None:
+ init_call_count[0] += 1
+ self.handle = handle
+ self.state = handle.getValueState("count", "int")
+
+ def handleInputRows(self, key, rows, timerValues) -> Iterator:
+ val = self.state.get() if self.state.exists() else None
+ count = val[0] if val else 0
+ for row in rows:
+ count += 1
+ self.state.update((count,))
+ yield Row(key=key[0], count=count)
+
+ processor = InitCountingProcessor()
+ tester = TwsTester(processor)
+ tester.test("key1", [Row(value="a")])
+ tester.test("key1", [Row(value="b")])
+ self.assertEqual(init_call_count[0], 1)
+
+ def test_processor_state_persists_across_batches(self):
Review Comment:
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r2755801161
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
Review Comment:
This test is for scenario when user tries to clear list state that does not exist. This is a corner case.
In tws_tester.py on line 132 there is a check that we remove state from the dict only if it exists. If we don't do that check, this test fails with KeyError.
I added 2 tests for clearing map and value 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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on code in PR #53758:
URL: https://github.com/apache/spark/pull/53758#discussion_r2755784817
##########
python/pyspark/sql/tests/pandas/streaming/test_tws_tester.py:
##########
@@ -0,0 +1,1294 @@
+#
+# 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 random
+import tempfile
+import unittest
+
+import pandas as pd
+import pandas.testing as pdt
+
+from pyspark import SparkConf
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import split
+from pyspark.sql.streaming import StatefulProcessor, TwsTester
+from pyspark.sql.streaming.query import StreamingQuery
+from pyspark.sql.tests.pandas.helper.helper_pandas_transform_with_state import (
+ AllMethodsTestProcessorFactory,
+ EventTimeCountProcessorFactory,
+ EventTimeSessionProcessorFactory,
+ RunningCountStatefulProcessorFactory,
+ SessionTimeoutProcessorFactory,
+ TopKProcessorFactory,
+ WordFrequencyProcessorFactory,
+)
+from pyspark.sql.types import (
+ DoubleType,
+ IntegerType,
+ Row,
+ StringType,
+ StructField,
+ StructType,
+)
+from pyspark.errors import PySparkValueError, PySparkAssertionError
+from pyspark.errors.exceptions.base import IllegalArgumentException
+from pyspark.testing.sqlutils import (
+ ReusedSQLTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message or "",
+)
+class TwsTesterTests(ReusedSQLTestCase):
+ def test_running_count_processor(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ self.assertEqual(tester.test("key1", [Row(value="a")]), [Row(key="key1", count=1)])
+ self.assertEqual(
+ tester.test("key2", [Row(value="a"), Row(value="a")]),
+ [Row(key="key2", count=2)],
+ )
+ self.assertEqual(tester.test("key3", [Row(value="a")]), [Row(key="key3", count=1)])
+ self.assertEqual(
+ tester.test("key1", [Row(value="a"), Row(value="a"), Row(value="a")]),
+ [Row(key="key1", count=4)],
+ )
+
+ self.assertEqual(tester.peekValueState("count", "key1"), (4,))
+ self.assertEqual(tester.peekValueState("count", "key2"), (2,))
+ self.assertEqual(tester.peekValueState("count", "key3"), (1,))
+ self.assertIsNone(tester.peekValueState("count", "key4"))
+
+ def test_running_count_processor_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"value": ["a", "a"]}))
+ expected2 = pd.DataFrame({"key": ["key2"], "count": [2]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"value": ["a"]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "count": [1]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"value": ["a", "a", "a"]}))
+ expected4 = pd.DataFrame({"key": ["key1"], "count": [4]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_value_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.updateValueState("count", "foo", (5,))
+ tester.test("foo", [Row(value="q")])
+ self.assertEqual(tester.peekValueState("count", "foo"), (6,))
+
+ def test_topk_processor(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test("key1", [Row(score=2.0), Row(score=3.0), Row(score=1.0)])
+ self.assertEqual(ans1, [Row(key="key1", score=3.0), Row(key="key1", score=2.0)])
+
+ ans2 = tester.test(
+ "key2", [Row(score=10.0), Row(score=20.0), Row(score=30.0), Row(score=40.0)]
+ )
+ self.assertEqual(ans2, [Row(key="key2", score=40.0), Row(key="key2", score=30.0)])
+
+ ans3 = tester.test("key3", [Row(score=100.0)])
+ self.assertEqual(ans3, [Row(key="key3", score=100.0)])
+
+ self.assertEqual(tester.peekListState("topK", "key1"), [(3.0,), (2.0,)])
+ self.assertEqual(tester.peekListState("topK", "key2"), [(40.0,), (30.0,)])
+ self.assertEqual(tester.peekListState("topK", "key3"), [(100.0,)])
+ self.assertEqual(tester.peekListState("topK", "key4"), [])
+
+ ans4 = tester.test("key1", [Row(score=10.0)])
+ self.assertEqual(ans4, [Row(key="key1", score=10.0), Row(key="key1", score=3.0)])
+ self.assertEqual(tester.peekListState("topK", "key1"), [(10.0,), (3.0,)])
+
+ def test_topk_processor_pandas(self):
+ processor = TopKProcessorFactory(k=2).pandas()
+ tester = TwsTester(processor)
+
+ ans1 = tester.testInPandas("key1", pd.DataFrame({"score": [2.0, 3.0, 1.0]}))
+ expected1 = pd.DataFrame({"key": ["key1", "key1"], "score": [3.0, 2.0]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("key2", pd.DataFrame({"score": [10.0, 20.0, 30.0, 40.0]}))
+ expected2 = pd.DataFrame({"key": ["key2", "key2"], "score": [40.0, 30.0]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ ans3 = tester.testInPandas("key3", pd.DataFrame({"score": [100.0]}))
+ expected3 = pd.DataFrame({"key": ["key3"], "score": [100.0]})
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ ans4 = tester.testInPandas("key1", pd.DataFrame({"score": [10.0]}))
+ expected4 = pd.DataFrame({"key": ["key1", "key1"], "score": [10.0, 3.0]})
+ pdt.assert_frame_equal(ans4, expected4, check_like=True)
+
+ def test_direct_access_to_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+
+ tester.updateListState("topK", "a", [(6.0,), (5.0,)])
+ tester.updateListState("topK", "b", [(8.0,), (7.0,)])
+ tester.test("a", [Row(score=10.0)])
+ tester.test("b", [Row(score=7.5)])
+ tester.test("c", [Row(score=1.0)])
+
+ assert tester.peekListState("topK", "a") == [(10.0,), (6.0,)]
+ assert tester.peekListState("topK", "b") == [(8.0,), (7.5,)]
+ assert tester.peekListState("topK", "c") == [(1.0,)]
+ assert tester.peekListState("topK", "d") == []
+
+ def test_word_frequency_processor(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ ans1 = tester.test(
+ "user1",
+ [
+ Row(word="hello"),
+ Row(word="world"),
+ Row(word="hello"),
+ Row(word="world"),
+ ],
+ )
+ self.assertEqual(
+ ans1,
+ [
+ Row(key="user1", word="hello", count=1),
+ Row(key="user1", word="world", count=1),
+ Row(key="user1", word="hello", count=2),
+ Row(key="user1", word="world", count=2),
+ ],
+ )
+
+ ans2 = tester.test("user2", [Row(word="hello"), Row(word="spark")])
+ self.assertEqual(
+ ans2,
+ [
+ Row(key="user2", word="hello", count=1),
+ Row(key="user2", word="spark", count=1),
+ ],
+ )
+
+ # Check state using peekMapState.
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (2,), ("world",): (2,)},
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user2"),
+ {("hello",): (1,), ("spark",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {})
+
+ # Process more data for user1.
+ ans3 = tester.test("user1", [Row(word="hello"), Row(word="test")])
+ self.assertEqual(
+ ans3,
+ [
+ Row(key="user1", word="hello", count=3),
+ Row(key="user1", word="test", count=1),
+ ],
+ )
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (3,), ("world",): (2,), ("test",): (1,)},
+ )
+
+ def test_word_frequency_processor_pandas(self):
+ processor = WordFrequencyProcessorFactory().pandas()
+ tester = TwsTester(processor)
+
+ input_df1 = pd.DataFrame({"word": ["hello", "world", "hello", "world"]})
+ ans1 = tester.testInPandas("user1", input_df1)
+ expected1 = pd.DataFrame(
+ {
+ "key": ["user1", "user1", "user1", "user1"],
+ "word": ["hello", "world", "hello", "world"],
+ "count": [1, 1, 2, 2],
+ }
+ )
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ input_df2 = pd.DataFrame({"word": ["hello", "spark"]})
+ ans2 = tester.testInPandas("user2", input_df2)
+ expected2 = pd.DataFrame(
+ {
+ "key": ["user2", "user2"],
+ "word": ["hello", "spark"],
+ "count": [1, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ input_df3 = pd.DataFrame({"word": ["hello", "test"]})
+ ans3 = tester.testInPandas("user1", input_df3)
+ expected3 = pd.DataFrame(
+ {
+ "key": ["user1", "user1"],
+ "word": ["hello", "test"],
+ "count": [3, 1],
+ }
+ )
+ pdt.assert_frame_equal(ans3, expected3, check_like=True)
+
+ def test_direct_access_to_map_state(self):
+ processor = WordFrequencyProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ tester.updateMapState("frequencies", "user1", {("hello",): (5,), ("world",): (3,)})
+ tester.updateMapState("frequencies", "user2", {("spark",): (10,)})
+
+ tester.test("user1", [Row(word="hello"), Row(word="goodbye")])
+ tester.test("user2", [Row(word="spark")])
+ tester.test("user3", [Row(word="new")])
+
+ self.assertEqual(
+ tester.peekMapState("frequencies", "user1"),
+ {("hello",): (6,), ("world",): (3,), ("goodbye",): (1,)},
+ )
+ self.assertEqual(tester.peekMapState("frequencies", "user2"), {("spark",): (11,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user3"), {("new",): (1,)})
+ self.assertEqual(tester.peekMapState("frequencies", "user4"), {})
+
+ # Example of how TwsTester can be used to test step function.
+ def test_step_function(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function using TwsTester to inspect how processing a single row changes
+ # state.
+ def step_function(key: str, input_row: str, state_in: int) -> int:
+ tester.updateValueState("count", key, (state_in,))
+ tester.test(key, [Row(value=input_row)])
+ return tester.peekValueState("count", key)[0]
+
+ self.assertEqual(step_function("key1", "a", 10), 11)
+
+ # Example of how TwsTester can be used to simulate real-time mode (row by row processing).
+ def test_row_by_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ # Example of helper function to test how TransformWithState processes rows one-by-one,
+ # which can be used to simulate real-time mode.
+ def test_row_by_row_helper(input_rows: list[tuple[str, Row]]) -> list[Row]:
+ result: list[Row] = []
+ for key, row in input_rows:
+ result += tester.test(key, [row])
+ return result
+
+ ans = test_row_by_row_helper(
+ [
+ ("key1", Row(value="a")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key2", Row(value="b")),
+ ("key1", Row(value="c")),
+ ("key1", Row(value="c")),
+ ("key3", Row(value="q")),
+ ]
+ )
+ self.assertEqual(
+ ans,
+ [
+ Row(key="key1", count=1),
+ Row(key="key2", count=1),
+ Row(key="key1", count=2),
+ Row(key="key2", count=2),
+ Row(key="key1", count=3),
+ Row(key="key1", count=4),
+ Row(key="key3", count=1),
+ ],
+ )
+
+ # Tests that TwsTester calls handleInitialState.
+ def test_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(
+ processor,
+ initialStateRow=[
+ ("a", Row(initial_count=10)),
+ ("b", Row(initial_count=20)),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.test("a", [Row(value="a")])
+ self.assertEqual(ans1, [Row(key="a", count=11)])
+
+ ans2 = tester.test("c", [Row(value="c")])
+ self.assertEqual(ans2, [Row(key="c", count=1)])
+
+ def test_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(
+ processor,
+ initialStatePandas=[
+ ("a", pd.DataFrame({"initial_count": [10]})),
+ ("b", pd.DataFrame({"initial_count": [20]})),
+ ],
+ )
+ self.assertEqual(tester.peekValueState("count", "a"), (10,))
+ self.assertEqual(tester.peekValueState("count", "b"), (20,))
+
+ ans1 = tester.testInPandas("a", pd.DataFrame({"value": ["a"]}))
+ expected1 = pd.DataFrame({"key": ["a"], "count": [11]})
+ pdt.assert_frame_equal(ans1, expected1, check_like=True)
+
+ ans2 = tester.testInPandas("c", pd.DataFrame({"value": ["c"]}))
+ expected2 = pd.DataFrame({"key": ["c"], "count": [1]})
+ pdt.assert_frame_equal(ans2, expected2, check_like=True)
+
+ def test_all_methods_processor(self):
+ """Test that TwsTester exercises all state methods."""
+ processor = AllMethodsTestProcessorFactory().row()
+ tester = TwsTester(processor)
+
+ results = tester.test(
+ "k",
+ [
+ Row(cmd="value-exists"), # false
+ Row(cmd="value-set"), # set to 42
+ Row(cmd="value-exists"), # true
+ Row(cmd="value-clear"), # clear
+ Row(cmd="value-exists"), # false again
+ Row(cmd="list-exists"), # false
+ Row(cmd="list-append"), # append a, b
+ Row(cmd="list-exists"), # true
+ Row(cmd="list-append-array"), # append c, d
+ Row(cmd="list-get"), # a,b,c,d
+ Row(cmd="map-exists"), # false
+ Row(cmd="map-add"), # add x=1, y=2, z=3
+ Row(cmd="map-exists"), # true
+ Row(cmd="map-keys"), # x,y,z
+ Row(cmd="map-values"), # 1,2,3
+ Row(cmd="map-iterator"), # x=1,y=2,z=3
+ Row(cmd="map-remove"), # remove y
+ Row(cmd="map-keys"), # x,z
+ Row(cmd="map-clear"), # clear map
+ Row(cmd="map-exists"), # false
+ ],
+ )
+
+ self.assertEqual(
+ results,
+ [
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="value-set:done"),
+ Row(key="k", result="value-exists:True"),
+ Row(key="k", result="value-clear:done"),
+ Row(key="k", result="value-exists:False"),
+ Row(key="k", result="list-exists:False"),
+ Row(key="k", result="list-append:done"),
+ Row(key="k", result="list-exists:True"),
+ Row(key="k", result="list-append-array:done"),
+ Row(key="k", result="list-get:a,b,c,d"),
+ Row(key="k", result="map-exists:False"),
+ Row(key="k", result="map-add:done"),
+ Row(key="k", result="map-exists:True"),
+ Row(key="k", result="map-keys:x,y,z"),
+ Row(key="k", result="map-values:1,2,3"),
+ Row(key="k", result="map-iterator:x=1,y=2,z=3"),
+ Row(key="k", result="map-remove:done"),
+ Row(key="k", result="map-keys:x,z"),
+ Row(key="k", result="map-clear:done"),
+ Row(key="k", result="map-exists:False"),
+ ],
+ )
+
+ def test_both_initial_states_specified(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ with self.assertRaises(AssertionError):
+ TwsTester(
+ processor,
+ initialStateRow=[("a", Row(initial_count=10))],
+ initialStatePandas=[("a", pd.DataFrame({"initial_count": [10]}))],
+ )
+
+ def test_timer_registration_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.registerTimer(12345)
+
+ def test_delete_timer_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ tester.handle.deleteTimer(12345)
+
+ def test_list_timers_raises_error_in_none_mode(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, timeMode="None")
+ with self.assertRaisesRegex(PySparkValueError, "UNSUPPORTED_OPERATION"):
+ list(tester.handle.listTimers())
+
+ def test_empty_input_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [])
+ self.assertEqual(result, [Row(key="key1", count=0)])
+
+ def test_empty_input_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": []})
+ result = tester.testInPandas("key1", input_df)
+ expected = pd.DataFrame({"key": ["key1"], "count": [0]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_empty_initial_state_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor, initialStateRow=[])
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_empty_initial_state_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor, initialStatePandas=[])
+ result = tester.testInPandas("key1", pd.DataFrame({"value": ["a"]}))
+ expected = pd.DataFrame({"key": ["key1"], "count": [1]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_none_values_in_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value=None), Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=2)])
+
+ def test_single_row_input(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key1", [Row(value="a")])
+ self.assertEqual(result, [Row(key="key1", count=1)])
+
+ def test_peek_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.peekValueState("count", "nonexistent_key")
+ self.assertIsNone(result)
+
+ def test_update_nonexistent_state(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.updateValueState("nonexistent_state", "key1", (5,))
+
+ def test_peek_state_before_initialization(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ with self.assertRaises(AssertionError):
+ tester.peekValueState("nonexistent_state", "key1")
+
+ def test_clear_nonexistent_list_state(self):
+ processor = TopKProcessorFactory(k=2).row()
+ tester = TwsTester(processor)
+ tester.handle.setGroupingKey("key1")
+ list_state = tester.handle.getListState("topK", "double")
+ list_state.clear()
+
+ def test_delete_if_exists_nonexistent(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ tester.handle.deleteIfExists("nonexistent_state")
+
+ def test_numeric_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ self.assertEqual(tester.test(1, [Row(value="a"), Row(value="c")]), [Row(key=1, count=2)])
+ self.assertEqual(tester.test(2, [Row(value="b")]), [Row(key=2, count=1)])
+
+ def test_tuple_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test((1, 2), [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=(1, 2), count=2)])
+
+ def test_special_characters_in_key(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test("key@#$%", [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key="key@#$%", count=2)])
+
+ def test_null_key_value_row(self):
+ processor = RunningCountStatefulProcessorFactory().row()
+ tester = TwsTester(processor)
+ result = tester.test(None, [Row(value="a"), Row(value="b")])
+ self.assertEqual(result, [Row(key=None, count=2)])
+
+ def test_null_key_value_pandas(self):
+ processor = RunningCountStatefulProcessorFactory().pandas()
+ tester = TwsTester(processor)
+ input_df = pd.DataFrame({"value": ["a", "b"]})
+ result = tester.testInPandas(None, input_df)
+ expected = pd.DataFrame({"key": [None], "count": [2]})
+ pdt.assert_frame_equal(result, expected, check_like=True)
+
+ def test_multiple_value_states(self):
Review Comment:
Removed these tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on PR #53758:
URL: https://github.com/apache/spark/pull/53758#issuecomment-3849201695
@HeartSaVioR , the error is failure to import pandas. I can try to fix it by importing it conditionally only when testInPandas is used. However, I don't understand why the same code was passing CI few weeks ago.
--
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-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "fedimser (via GitHub)" <gi...@apache.org>.
fedimser commented on PR #53758:
URL: https://github.com/apache/spark/pull/53758#issuecomment-3850551474
@HeartSaVioR , in previous CI run all jobs passed except pyspark-connect. And in latest run (after rebase), pyspark-connect passed but other tests failed. Is this sufficient signal that this PR doesn't break any tests?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-54805][SS][PYTHON] Implement TwsTester in PySpark [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #53758:
URL: https://github.com/apache/spark/pull/53758#issuecomment-3850960014
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