spark-test commented on code in PR #59776:
URL: https://github.com/apache/spark/pull/59776#discussion_r4372946180
##########
python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py:
##########
@@ -272,107 +278,121 @@ def run(self, split_index, data):
class HandlerUDFsTests(PySparkBaseTestCase):
class _Handler(BatchEvalTypeHandler["pa.RecordBatch"]):
+ _multi_udf = True
+
def run(self, split_index, data):
return data
@staticmethod
def _stop(x):
raise StopIteration(x)
- def _new(self, handler_cls, runner_conf=None):
- return handler_cls(
- [
- UDF(
- raw_func=self._stop,
- args_offsets=[1],
- kwargs_offsets={"y": 2},
- return_type=LongType(),
- result_id=0,
- )
- ],
- runner_conf or RunnerConf({}),
- None,
+ @staticmethod
+ def _udf(*funcs):
+ return UDF(
+ funcs=list(funcs),
+ args_offsets=[1],
+ kwargs_offsets={"y": 2},
+ return_type=LongType(),
+ result_id=0,
)
+ def _new(self, handler_cls, runner_conf=None):
+ return handler_cls([self._udf(self._stop)], runner_conf or
RunnerConf({}), None)
+
def test_wraps_udfs(self):
handler = self._new(self._Handler)
- udf = handler._udfs[0]
+ udf = handler._udf_descs[0]
self.assertIsInstance(udf, UDF)
self.assertEqual(udf.args_offsets, [1])
self.assertEqual(udf.kwargs_offsets, {"y": 2})
self.assertEqual(udf.return_type, LongType())
- # The handler keeps the function wrapped in fail_on_stopiteration; the
record keeps it
- # unwrapped.
+ # The handler keeps the function wrapped in fail_on_stopiteration; the
record and the
+ # handler's raw function keep it unwrapped.
with self.assertRaises(PySparkRuntimeError):
handler._funcs[0](1)
with self.assertRaises(StopIteration):
- udf.raw_func(1)
-
- def test_fail_on_stopiteration_per_handler(self):
- not_wrapped = {
- PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
- PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
- PythonEvalType.SQL_ARROW_BATCHED_UDF,
- }
- for _, variants in list(_eval_type_handlers.values()):
- for handler_cls in variants.values():
- with self.subTest(handler=handler_cls.__name__):
- self.assertEqual(
- handler_cls._fail_on_stopiteration,
- handler_cls.eval_type not in not_wrapped,
- )
+ udf.funcs[0](1)
+ self.assertIs(handler._raw_funcs[0], self._stop)
- class _NotWrapped(self._Handler):
- _fail_on_stopiteration = False
-
- with self.assertRaises(StopIteration):
- self._new(_NotWrapped)._funcs[0](1)
+ def test_chains_funcs_in_order(self):
+ handler = self._Handler(
+ [self._udf(lambda x: x + 1, lambda x: x * 2, str)],
RunnerConf({}), None
+ )
+ self.assertEqual(handler._raw_funcs[0](3), "8")
+ self.assertEqual(handler._funcs[0](3), "8")
- def test_profiler_selected_by_conf(self):
- def f(x):
- return x
+ def test_single_func_not_chained(self):
+ # A single function is kept as is, so its signature (and an
incremental Aggregator
+ # object) survives.
+ def f(key, value):
+ return value
- def data_source_f(x):
- return x
+ handler = self._Handler([self._udf(f)], RunnerConf({}), None)
+ self.assertIs(handler._raw_funcs[0], f)
- data_source_f.__module__ = "pyspark.sql.worker.plan_data_source_read"
+ def test_single_udf(self):
+ class _Single(BatchEvalTypeHandler["pa.RecordBatch"]):
+ def run(self, split_index, data):
+ return data
- udf_perf = self._Handler([],
RunnerConf({"spark.sql.pyspark.udf.profiler": "perf"}), None)
- data_source_perf = self._Handler(
- [], RunnerConf({"spark.sql.pyspark.dataSource.profiler": "perf"}),
None
- )
- udf_memory = self._Handler(
- [], RunnerConf({"spark.sql.pyspark.udf.profiler": "memory"}), None
- )
- no_profiler = self._Handler([], RunnerConf({}), None)
- with patch.object(
- EvalTypeHandler, "_with_perf_profiler", side_effect=lambda func,
result_id: "perf"
- ) as perf:
- self.assertEqual(udf_perf._with_profiler(f, 0), "perf")
- self.assertIs(data_source_perf._with_profiler(f, 0), f)
- self.assertIs(udf_perf._with_profiler(data_source_f, 0),
data_source_f)
- self.assertEqual(data_source_perf._with_profiler(data_source_f,
0), "perf")
- self.assertEqual(perf.call_count, 2)
- with patch.object(
- EvalTypeHandler, "_with_memory_profiler", side_effect=lambda func,
result_id: "memory"
+ def f(x):
+ return x + 1
+
+ handler = _Single([self._udf(f)], RunnerConf({}), None)
+ self.assertIs(handler._udf_desc, handler._udf_descs[0])
+ self.assertIs(handler._raw_func, f)
+ self.assertIs(handler._func, handler._funcs[0])
+ self.assertEqual(handler._func(1), 2)
+ for udf_descs in ([], [self._udf(f), self._udf(f)]):
+ with self.subTest(num_udfs=len(udf_descs)):
+ with self.assertRaisesRegex(
+ PySparkRuntimeError,
+ "^One _Single UDF expected, but got
{}.$".format(len(udf_descs)),
+ ):
+ _Single(udf_descs, RunnerConf({}), None)
+
+ handler = self._Handler([self._udf(f), self._udf(f)], RunnerConf({}),
None)
+ self.assertFalse(hasattr(handler, "_udf_desc"))
+
+ def test_single_udf_count_error_names_eval_type(self):
+ for eval_type, expected in (
+ (PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF, "COGROUPED_MAP_ARROW
UDF"),
+ (
+ PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE,
+ "GROUPED_MAP_PANDAS_UDF_WITH_STATE UDF",
+ ),
+ (PythonEvalType.SQL_TABLE_UDF, "UDTF"),
):
- self.assertEqual(udf_memory._with_profiler(f, 0), "memory")
- self.assertIs(no_profiler._with_profiler(f, 0), f)
-
-
-class UnpickleAndChainTests(PySparkBaseTestCase):
- def test_chains_in_order(self):
- ser = CloudPickleSerializer()
- func, return_type = unpickle_and_chain(
- pickleSer,
- [
- ser.dumps((lambda x: x + 1, LongType())),
- ser.dumps((lambda x: str(x * 2), StringType())),
- ],
- )
- self.assertEqual(func(3), "8")
- self.assertEqual(return_type, StringType())
+ handler_cls = _eval_type_handlers[eval_type][1][None]
+ with self.subTest(handler=handler_cls.__name__):
+ with self.assertRaisesRegex(
+ PySparkRuntimeError, "^One {} expected, but got
2.$".format(expected)
+ ):
+ handler_cls([self._udf(str), self._udf(str)],
RunnerConf({}), None)
Review Comment:
**Nit (P3):** `ArrowCoGroupedMapUDFHandler.__init__` and
`PandasGroupedMapWithStateUDFHandler.__init__` call
`require_minimum_pyarrow_version()` (the latter also
`require_minimum_pandas_version()`) before `super().__init__` reaches the
UDF-count check. Without those optional packages, this call raises
`PySparkImportError`, which subclasses `ImportError` rather than
`PySparkRuntimeError`. So `assertRaisesRegex` lets it through and the two
subtests error instead of skipping.
I reproduced this in an environment without pandas and pyarrow: the
`COGROUPED_MAP_ARROW` and `GROUPED_MAP_PANDAS_UDF_WITH_STATE` subtests error,
while the UDTF subtest and the rest of the module pass. Before this PR, the
module built only test-local handlers and ran cleanly without the sql extras.
The sibling eval_handlers test modules guard such tests with `have_pandas` /
`have_pyarrow` skips, and the same guard here would make the test skip. CI
images ship both packages, so this shows up only in local or downstream runs
without them.
**Verification:**
- **Compatibility:** In an environment without pyarrow (and one without
pandas), the base eval-handler tests finish with the naming test skipped and no
errors; with both installed, the naming assertions still run and pass.
##########
python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py:
##########
@@ -272,107 +278,121 @@ def run(self, split_index, data):
class HandlerUDFsTests(PySparkBaseTestCase):
class _Handler(BatchEvalTypeHandler["pa.RecordBatch"]):
+ _multi_udf = True
+
def run(self, split_index, data):
return data
@staticmethod
def _stop(x):
raise StopIteration(x)
- def _new(self, handler_cls, runner_conf=None):
- return handler_cls(
- [
- UDF(
- raw_func=self._stop,
- args_offsets=[1],
- kwargs_offsets={"y": 2},
- return_type=LongType(),
- result_id=0,
- )
- ],
- runner_conf or RunnerConf({}),
- None,
+ @staticmethod
+ def _udf(*funcs):
+ return UDF(
+ funcs=list(funcs),
+ args_offsets=[1],
+ kwargs_offsets={"y": 2},
+ return_type=LongType(),
+ result_id=0,
)
+ def _new(self, handler_cls, runner_conf=None):
+ return handler_cls([self._udf(self._stop)], runner_conf or
RunnerConf({}), None)
+
def test_wraps_udfs(self):
handler = self._new(self._Handler)
- udf = handler._udfs[0]
+ udf = handler._udf_descs[0]
self.assertIsInstance(udf, UDF)
self.assertEqual(udf.args_offsets, [1])
self.assertEqual(udf.kwargs_offsets, {"y": 2})
self.assertEqual(udf.return_type, LongType())
- # The handler keeps the function wrapped in fail_on_stopiteration; the
record keeps it
- # unwrapped.
+ # The handler keeps the function wrapped in fail_on_stopiteration; the
record and the
+ # handler's raw function keep it unwrapped.
with self.assertRaises(PySparkRuntimeError):
handler._funcs[0](1)
with self.assertRaises(StopIteration):
- udf.raw_func(1)
-
- def test_fail_on_stopiteration_per_handler(self):
Review Comment:
**Non-blocking (P2):** This hunk drops
`test_fail_on_stopiteration_per_handler` and `test_profiler_selected_by_conf`,
and nothing in head replaces them, even though `EvalTypeHandler.__init__` is
rewritten around the loop that applies both mechanisms.
`test_multi_udf_per_handler` pins the new flag, but the
`_fail_on_stopiteration` table it replaced is no longer checked.
What loses its unit-level failure signal:
- Of the 30 registered UDF handlers that wrap in `fail_on_stopiteration` by
default, only `ArrowElementwiseUDFHandler` has a handler-level StopIteration
test. The end-to-end `test_stopiteration_in_*` tests only match the regex
`StopIteration`, which a PEP 479 `generator raised StopIteration` RuntimeError
also satisfies. So setting `_fail_on_stopiteration = False` on, say,
`ArrowGroupedMapUDFHandler` would pass the unit suite, losing the guard that
turns a StopIteration raised in user code into `STOP_ITERATION_OCCURRED`
instead of silently ending the loop.
- `_with_profiler` routing: the UDF profiler and the data-source profiler
isolated in both directions, memory selection, and returning the function
unchanged when no profiler is set.
Both tests still work against the head fixture (`_Handler` now sets
`_multi_udf = True`, so `_Handler([], ...)` and `_new(...)` stay valid). I ran
their merge-base bodies against this head and both pass, so restoring them next
to the new tests is enough. Runtime behavior is not affected; this is lost
regression coverage only.
**Verification:**
- **Regression:** Coverage fails when any registered handler's
_fail_on_stopiteration differs from the intended opt-out set, or when a handler
with the flag off still wraps its function in fail_on_stopiteration.
- **Regression:** Coverage fails when profiler selection routes a
data-source function to the UDF profiler or a regular function to the
data-source profiler, picks the wrong profiler for perf or memory, or wraps the
function when no profiler is configured.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
