Yicong-Huang commented on code in PR #59776:
URL: https://github.com/apache/spark/pull/59776#discussion_r4381889454
##########
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:
Restored both tests unchanged in 96da283f461
(`test_base_eval_type_handlers.py:326`, `:346`); they were dropped by mistake.
##########
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:
Added a `have_pandas`/`have_pyarrow` skip to the test in 96da283f461
(`test_base_eval_type_handlers.py:417`), matching the sibling modules.
--
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]
