This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 5cb18010f352 [MINOR][PYTHON][TEST] Add TaskContext tests for scalar
Arrow UDFs
5cb18010f352 is described below
commit 5cb18010f35253b9e133b3e832be6cd0b127df60
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Tue Aug 4 17:32:25 2026 +0800
[MINOR][PYTHON][TEST] Add TaskContext tests for scalar Arrow UDFs
### What changes were proposed in this pull request?
Add a test to ScalarArrowUDFTestsMixin that verifies TaskContext is
initialized in both scalar and scalar-iterator Arrow UDFs. The test checks the
partition IDs returned by TaskContext and is also inherited by the Spark
Connect parity suite.
### Why are the changes needed?
TaskContext is covered for Pandas UDFs and mapInArrow, but scalar Arrow
UDFs did not have equivalent coverage. This test closes that gap and guards
classic and Spark Connect execution against regressions.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
build/sbt -Phive package
conda run -n spark-dev-313 --no-capture-output python/run-tests
--testnames "pyspark.sql.tests.arrow.test_arrow_udf_scalar
ScalarArrowUDFTests.test_arrow_udf_task_context"
conda run -n spark-dev-313 --no-capture-output python/run-tests
--testnames "pyspark.sql.tests.connect.arrow.test_parity_arrow_udf_scalar
ScalarArrowPythonUDFParityTests.test_arrow_udf_task_context"
conda run -n spark-dev-313 ruff check
python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py
conda run -n spark-dev-313 ruff format --check
python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #57744 from zhengruifeng/task-context-arrow-udf-dev1.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../sql/tests/arrow/test_arrow_udf_scalar.py | 27 ++++++++++++++++++++++
1 file changed, 27 insertions(+)
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py
b/python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py
index d574c7266537..651e4b94a4d0 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_udf_scalar.py
@@ -1267,6 +1267,33 @@ class ScalarArrowUDFTestsMixin:
result = df.select(return_one("id").alias("one")).collect()
self.assertEqual(expected, result)
+ def test_arrow_udf_task_context(self):
+ import pyarrow as pa
+
+ from pyspark import TaskContext
+
+ def partition_ids(size):
+ task_context = TaskContext.get()
+ assert task_context is not None
+ return pa.array([task_context.partitionId()] * size,
type=pa.int32())
+
+ @arrow_udf("int")
+ def scalar_task_context(values):
+ return partition_ids(len(values))
+
+ @arrow_udf("int", ArrowUDFType.SCALAR_ITER)
+ def scalar_iter_task_context(iterator):
+ for values in iterator:
+ yield partition_ids(len(values))
+
+ for task_context_udf in [scalar_task_context,
scalar_iter_task_context]:
+ rows = (
+ self.spark.range(10, numPartitions=2)
+ .select(task_context_udf("id").alias("partition_id"))
+ .collect()
+ )
+ self.assertEqual({0, 1}, {row.partition_id for row in rows})
+
class ScalarArrowUDFTests(ScalarArrowUDFTestsMixin, ReusedSQLTestCase):
@classmethod
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
