This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new be855c94fecc [SPARK-57677][PYTHON][TEST] Make UDTF cleanup tests
tolerate task retries
be855c94fecc is described below
commit be855c94fecc43030fada62e9a65b01a3ec50a75
Author: Tian Gao <[email protected]>
AuthorDate: Fri Jun 26 11:33:25 2026 -0700
[SPARK-57677][PYTHON][TEST] Make UDTF cleanup tests tolerate task retries
### What changes were proposed in this pull request?
The two UDTF cleanup tests `test_udtf_cleanup_with_exception_in_eval` and
`test_udtf_cleanup_with_exception_in_terminate` asserted that the file written
by the UDTF's `cleanup` method contained exactly the string `"cleanup"`. This
relaxes the assertions to check that `"cleanup"` is present (and, for the eval
test, that `"terminate"` is not), instead of requiring an exact match.
### Why are the changes needed?
A UDTF task may be retried by the scheduler, in which case `cleanup` runs
more than once and the file ends up with multiple `"cleanup"` strings. The
exact-match assertion would then fail. Checking for presence makes the tests
robust to retries.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
CI.
### Was this patch authored or co-authored using generative AI tooling?
Yes, Claude Code (Opus 4.8 high)
Closes #56755 from gaogaotiantian/SPARK-57677-udtf-cleanup.
Authored-by: Tian Gao <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
(cherry picked from commit 6fa1eab00b63e360b239c4399cf6866f6f394b0a)
Signed-off-by: Dongjoon Hyun <[email protected]>
---
python/pyspark/sql/tests/test_udtf.py | 12 +++++++++---
1 file changed, 9 insertions(+), 3 deletions(-)
diff --git a/python/pyspark/sql/tests/test_udtf.py
b/python/pyspark/sql/tests/test_udtf.py
index d623069eae6a..a272cee4c7b5 100644
--- a/python/pyspark/sql/tests/test_udtf.py
+++ b/python/pyspark/sql/tests/test_udtf.py
@@ -565,8 +565,11 @@ class BaseUDTFTestsMixin:
with open(path, "r") as f:
data = f.read()
- # Only cleanup method should be called.
- self.assertEqual(data, "cleanup")
+ # Only the cleanup method should be called, not terminate. The
UDTF may be retried,
+ # so cleanup can run more than once and the file may contain
multiple "cleanup"
+ # strings; just check that "cleanup" is present and "terminate" is
not.
+ self.assertIn("cleanup", data)
+ self.assertNotIn("terminate", data)
def test_udtf_cleanup_with_exception_in_terminate(self):
with tempfile.TemporaryDirectory(
@@ -595,7 +598,10 @@ class BaseUDTFTestsMixin:
with open(path, "r") as f:
data = f.read()
- self.assertEqual(data, "cleanup")
+ # The cleanup method should be called even when terminate raises.
The UDTF may be
+ # retried, so cleanup can run more than once and the file may
contain multiple
+ # "cleanup" strings; just check that "cleanup" is present.
+ self.assertIn("cleanup", data)
def test_init_with_exception(self):
@udtf(returnType="x: int")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
