This is an automated email from the ASF dual-hosted git repository.
zhengruifeng 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 361b5ca79872 [SPARK-58859][PYTHON][TEST] Add BarrierTaskContext
negative tests for SQL map APIs
361b5ca79872 is described below
commit 361b5ca79872fa4f62a175f7c146ba5aa346e14c
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Aug 19 18:43:37 2026 +0800
[SPARK-58859][PYTHON][TEST] Add BarrierTaskContext negative tests for SQL
map APIs
### What changes were proposed in this pull request?
This PR strengthens the existing `mapInPandas` and `mapInArrow`
barrier-mode tests.
For the non-barrier case, the tests now call `BarrierTaskContext.get()`
directly and assert
that it fails. The existing checks continue to verify that
`TaskContext.get()` is present and
is not a `BarrierTaskContext` outside barrier mode.
### Why are the changes needed?
`BarrierTaskContext` should only be available in barrier execution. The
existing SQL map API
tests checked the `TaskContext.get()` type in both modes, but they did not
directly verify that
`BarrierTaskContext.get()` fails when `barrier=False`.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added negative coverage to existing PySpark SQL tests:
- `MapInPandasTestsMixin.test_map_in_pandas_with_barrier_mode`
- `MapInArrowTestsMixin.test_map_in_arrow_with_barrier_mode`
Local checks:
- `git diff --check`
- ASCII and line-length scans on changed files
Focused PySpark suites were not run locally.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58105 from zhengruifeng/SPARK-58859-barrier-context-negative-tests.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 1938bcf2451a27346685519f7d9a01046bad3447)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/sql/tests/arrow/test_arrow_map.py | 10 ++++++++++
python/pyspark/sql/tests/pandas/test_pandas_map.py | 10 ++++++++++
2 files changed, 20 insertions(+)
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_map.py
b/python/pyspark/sql/tests/arrow/test_arrow_map.py
index 5119e0e827f6..c8f6d9356aba 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_map.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_map.py
@@ -185,6 +185,16 @@ class MapInArrowTestsMixin:
def test_map_in_arrow_with_barrier_mode(self):
df = self.spark.range(10)
+ def func0(iterator):
+ from pyspark import BarrierTaskContext
+
+ BarrierTaskContext.get()
+ for batch in iterator:
+ yield batch
+
+ with self.assertRaisesRegex(PythonException,
"\\[NOT_IN_BARRIER_STAGE\\]"):
+ df.mapInArrow(func0, "id long", False).collect()
+
def func1(iterator):
from pyspark import TaskContext, BarrierTaskContext
diff --git a/python/pyspark/sql/tests/pandas/test_pandas_map.py
b/python/pyspark/sql/tests/pandas/test_pandas_map.py
index bfcedc6c8899..96d354e39fd7 100644
--- a/python/pyspark/sql/tests/pandas/test_pandas_map.py
+++ b/python/pyspark/sql/tests/pandas/test_pandas_map.py
@@ -459,6 +459,16 @@ class MapInPandasTestsMixin:
def test_map_in_pandas_with_barrier_mode(self):
df = self.spark.range(10)
+ def func0(iterator):
+ from pyspark import BarrierTaskContext
+
+ BarrierTaskContext.get()
+ for batch in iterator:
+ yield batch
+
+ with self.assertRaisesRegex(PythonException,
"\\[NOT_IN_BARRIER_STAGE\\]"):
+ df.mapInPandas(func0, "id long", False).collect()
+
def func1(iterator):
from pyspark import TaskContext, BarrierTaskContext
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
