gaogaotiantian commented on code in PR #59845:
URL: https://github.com/apache/spark/pull/59845#discussion_r4474671422
##########
python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py:
##########
@@ -367,6 +367,19 @@ def test_concurrency_level(self):
out = list(handler.run(0, iter([self._batch(a=list(range(10)), b=[0] *
10)])))
self.assertEqual(out[0].column(0).to_pylist(), [v * 2 for v in
range(10)])
+ def test_concurrency_level_multi_and_zero_arg(self):
+ conf =
RunnerConf({"spark.sql.execution.pythonUDF.arrow.concurrency.level": "4"})
+ handler = self._handler((lambda a, b: a - b, [0, 1], {}, LongType()),
runner_conf=conf)
+ out = list(handler.run(0, iter([self._batch(a=list(range(10)), b=[1] *
10)])))
+ self.assertEqual(out[0].column(0).to_pylist(), [v - 1 for v in
range(10)])
+
+ handler = self._handler(
+ (lambda: 7, [], {}, LongType()), runner_conf=conf,
input_type=StructType()
+ )
+ empty = pa.RecordBatch.from_pydict({"a": [1, 2, 3]}).select([])
+ out = list(handler.run(0, iter([empty])))
+ self.assertEqual(out[0].column(0).to_pylist(), [7, 7, 7])
+
def test_stop_iteration_raises_runtime_error(self):
Review Comment:
Thanks, good catch. The test now asserts `STOP_ITERATION_OCCURRED` for
zero-arg, single-arg (`args[0]`), multi-arg (`zip(*args)`), and keyword-arg
UDFs. It passes on the base commit and on the latest commit, and fails on
cdca21c75a5 with `'RESULT_ROWS_MISMATCH' != 'STOP_ITERATION_OCCURRED'` for the
single-, multi-, and keyword-arg cases.
While doing this I found that the thread-pool path (`arrow.concurrency.level
> 0`) is not consistent here, before or after this PR: a UDF that raises
`StopIteration` fails with a bare `RuntimeError: generator raised
StopIteration` instead of `STOP_ITERATION_OCCURRED`. This handler sets
`_fail_on_stopiteration = False` and wraps the whole batch evaluation, so the
`StopIteration` is re-raised inside `Executor.map`'s result generator, and PEP
479 turns it into a `RuntimeError` before `fail_on_stopiteration` sees it. The
task still fails, so no data is lost. The elementwise handler wraps each call
and is consistent at both levels. I'd like to keep this PR free of behavior
changes, so the test only covers the default path, and I'll fix the thread-pool
path in a separate follow-up.
##########
python/pyspark/eval_handlers/_arrow.py:
##########
@@ -879,13 +879,21 @@ def run(self, split_index: int, data:
Iterator[pa.RecordBatch]) -> Iterator[pa.R
phase_timers = None
@fail_on_stopiteration
- def _evaluate_batch_udf(udf_func: Any, rows: list[tuple]) -> list[Any]:
+ def _evaluate_batch_udf(udf_func: Any, args: list[list[Any]],
num_rows: int) -> list[Any]:
if runner_conf.arrow_concurrency_level <= 0:
- return [udf_func(*row) for row in rows]
+ if not args:
+ return [udf_func() for _ in range(num_rows)]
+ # Call the UDF in the comprehension body, not via
list(map(...)): map()
+ # would swallow a StopIteration raised by the UDF as end of
input.
+ if len(args) == 1:
+ return [udf_func(v) for v in args[0]]
+ return [udf_func(*row) for row in zip(*args)]
Review Comment:
Agreed, and thanks for the micro-benchmark. I updated the title and
description to match the final code. Single-arg UDFs build no per-row tuple.
Multi-arg UDFs still get one short-lived tuple per row from `zip`; only the
materialized list of row tuples is gone. Keyword-arg UDFs still build a tuple
and a dict per call in `wrap_kwargs_support`. I also re-ran
`ArrowBatchedUDFTimeBench` against the latest commit and replaced the table. As
you found, the `concat3_udf` gains are smaller than with the `map()` version.
##########
python/pyspark/eval_handlers/_arrow.py:
##########
@@ -879,13 +879,21 @@ def run(self, split_index: int, data:
Iterator[pa.RecordBatch]) -> Iterator[pa.R
phase_timers = None
@fail_on_stopiteration
- def _evaluate_batch_udf(udf_func: Any, rows: list[tuple]) -> list[Any]:
+ def _evaluate_batch_udf(udf_func: Any, args: list[list[Any]],
num_rows: int) -> list[Any]:
if runner_conf.arrow_concurrency_level <= 0:
- return [udf_func(*row) for row in rows]
+ if not args:
+ return [udf_func() for _ in range(num_rows)]
+ # Call the UDF in the comprehension body, not via
list(map(...)): map()
+ # would swallow a StopIteration raised by the UDF as end of
input.
+ if len(args) == 1:
+ return [udf_func(v) for v in args[0]]
+ return [udf_func(*row) for row in zip(*args)]
from concurrent.futures import ThreadPoolExecutor
with
ThreadPoolExecutor(max_workers=runner_conf.arrow_concurrency_level) as pool:
- return list(pool.map(lambda row: udf_func(*row), rows))
+ if not args:
+ return list(pool.map(lambda _: udf_func(),
range(num_rows)))
+ return list(pool.map(udf_func, *args))
Review Comment:
You're right. I removed that sentence; the description now says the
thread-pool change is only a simplification and per-row allocations there are
unchanged. Reusing one executor per task sounds good as a follow-up.
##########
python/pyspark/eval_handlers/_arrow.py:
##########
@@ -920,15 +928,11 @@ def _evaluate_batch_udf(udf_func: Any, rows: list[tuple])
-> list[Any]:
# Reuse the batch start for the first UDF, avoiding
another clock read.
# Later UDFs start input after the loop transition.
input_start_ns = clock()
- rows = (
- [() for _ in range(num_rows)]
- if zero_arg
- else list(zip(*[columns[o] for o in offsets]))
- )
+ args = [] if zero_arg else [columns[o] for o in offsets]
if phase_timers is not None:
udf_start_ns = clock()
input_duration_ns += udf_start_ns - input_start_ns
- results = _evaluate_batch_udf(udf_func, rows)
+ results = _evaluate_batch_udf(udf_func, args, num_rows)
Review Comment:
Good point. I mentioned the shift from `pythonInputPreparationTime` to
`pythonUDFExecutionTime` for multi-arg UDFs in the description.
##########
python/pyspark/eval_handlers/_arrow.py:
##########
@@ -920,15 +928,11 @@ def _evaluate_batch_udf(udf_func: Any, rows: list[tuple])
-> list[Any]:
# Reuse the batch start for the first UDF, avoiding
another clock read.
# Later UDFs start input after the loop transition.
input_start_ns = clock()
- rows = (
- [() for _ in range(num_rows)]
- if zero_arg
- else list(zip(*[columns[o] for o in offsets]))
- )
+ args = [] if zero_arg else [columns[o] for o in offsets]
Review Comment:
Done. `args_kwargs_offsets` stays empty for zero-arg UDFs, and the
`zero_arg` element of `_udf_infos` and the `_NoValue` placeholder column are
removed.
##########
python/pyspark/eval_handlers/_arrow.py:
##########
@@ -879,13 +879,21 @@ def run(self, split_index: int, data:
Iterator[pa.RecordBatch]) -> Iterator[pa.R
phase_timers = None
@fail_on_stopiteration
- def _evaluate_batch_udf(udf_func: Any, rows: list[tuple]) -> list[Any]:
+ def _evaluate_batch_udf(udf_func: Any, args: list[list[Any]],
num_rows: int) -> list[Any]:
Review Comment:
Agreed it would be good to keep these consistent. I'll leave it for a
follow-up to keep this PR focused.
--
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]
