dongjoon-hyun commented on code in PR #59845:
URL: https://github.com/apache/spark/pull/59845#discussion_r4464058668
##########
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:
`Executor.map()` zips the iterables internally (`fs = [self.submit(fn,
*args) for args in zip(*iterables)]`) and creates a `Future` per row, so this
change does not reduce per-row allocations on the thread-pool path. The
sentence "The `arrow.concurrency.level` thread-pool path changes the same way"
in the description may be misleading.
Separately (pre-existing, not for this PR): a new `ThreadPoolExecutor` is
created and shut down per UDF per batch. Reusing a single executor per task
could be a 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:
The PR title and description still describe the first commit: `return
list(map(udf_func, *args))` and "no tuple is built per row". With this
follow-up, multi-arg UDFs still create one tuple per row via `zip(*args)` (only
the materialized list of tuples is avoided), and keyword-argument UDFs also
build a tuple and a dict per call inside `wrap_kwargs_support`.
In a local micro-benchmark (3-arg UDF, 10k rows x 50 iterations, min of 7
runs, CPython 3.12), I got:
| Version | Time |
|---|---|
| Before this PR | 0.0647s |
| This PR (d13ff764dcc) | 0.0572s |
| `list(map(...))` (cdca21c75a5) | 0.0494s |
So the `concat3_udf` numbers in the description seem to be measured with the
`map()` version and may overstate the gain of the final code. Could you update
the title/description and re-run `ArrowBatchedUDFTimeBench` with the latest
commit?
##########
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:
Not a blocker: `ArrowBatchedUDFLegacyPandasHandler` still materializes
`list(zip(...))` per UDF per batch (`_pandas.py` L799-806), and
`_evaluate_elementwise_udf` in `_elementwise.py` (L243-249) is almost the same
helper as this one. If reducing per-row allocations is the goal, it may be
worth applying the same change to the legacy path, or extracting a shared
helper (with the `StopIteration` note and the zero-arg branch) so that the
three handlers do not drift apart. A follow-up is fine.
##########
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):
Review Comment:
nit: Since this test covers the new thread-pool branches, could you also add
a keyword-argument UDF here (like `test_passes_arg_by_keyword_offset`)? Then
every argument shape that goes through `pool.map(udf_func, *args)` is exercised.
##########
python/benchmarks/bench_eval_type.py:
##########
@@ -540,6 +540,7 @@ def _build_scenario(cls, name):
"identity_udf": (lambda x: x, None, [0]),
"stringify_udf": (lambda x: str(x), StringType(), [0]),
"nullcheck_udf": (lambda x: x is not None, BooleanType(), [0]),
+ "concat3_udf": (lambda a, b, c: f"{a}{b}{c}", StringType(), [0, 1, 2]),
Review Comment:
nit: `_udfs` is shared via `_ArrowBatchedBenchMixin`, so this also adds 8
scenarios to `ArrowBatchedUDFPeakmemBench` (24 -> 32 cases for each of the Time
and Peakmem benchmarks). If multi-arg peak memory is not a target, you may want
to add it only to `ArrowBatchedUDFTimeBench`. Also, the comment above
(`arg_offsets=[0] means ...`) no longer applies to all entries.
##########
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:
This test still checks only the exception type. `RESULT_ROWS_MISMATCH` is
also a `PySparkRuntimeError`, so the `list(map(...))` regression fixed by
d13ff764dcc is not caught here. I confirmed that all tests in
`ArrowBatchedUDFArrowHandlerTests` pass on the first commit (cdca21c75a5) as
well.
Could you assert the error condition, as
`test_elementwise_eval_type_handlers.py` does, and add a multi-arg UDF so that
both the `args[0]` and `zip(*args)` branches are covered?
```python
with self.assertRaises(PySparkRuntimeError) as ctx:
list(handler.run(0, iter([self._batch(a=[1], b=[0])])))
self.assertEqual(ctx.exception.getCondition(), "STOP_ITERATION_OCCURRED")
```
##########
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:
nit (optional): Since `_evaluate_batch_udf` now checks `if not args`, the
`zero_arg` flag looks redundant. If `args_kwargs_offsets` were kept empty
instead of `args_kwargs_offsets or (0,)` (L833), this line could be `args =
[columns[o] for o in offsets]`, and the `zero_arg` element of `_udf_infos`
(L829, L834) and the `_NoValue` placeholder column (L922) could go away.
##########
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:
nit: For multi-arg UDFs, the per-row tuple creation (`zip`) moved from
before `udf_start_ns` (input preparation) into `_evaluate_batch_udf` (UDF
execution). When detailed timing is enabled, part of the time moves from
`pythonInputPreparationTime` to `pythonUDFExecutionTime`. It is minor, but it
might be worth mentioning in the description since the per-phase metrics change
slightly.
--
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]
