Spenserrrr commented on code in PR #59246:
URL: https://github.com/apache/spark/pull/59246#discussion_r4251448633
##########
python/pyspark/worker.py:
##########
@@ -2382,45 +2384,72 @@ def _evaluate_batch_udf(udf_func, rows):
return list(pool.map(lambda row: udf_func(*row), rows))
def func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
- for input_batch in data:
- num_rows = input_batch.num_rows
-
- # --- Input: Arrow -> Python columns ---
- columns = [
- (
- [conv(v) for v in
ArrowTableToRowsConversion._to_pylist(col)]
- if conv is not None
- else ArrowTableToRowsConversion._to_pylist(col)
- )
- for col, conv in zip(input_batch.itercolumns(),
arrow_to_py_converters)
- ]
- if not columns:
- columns = [[_NoValue] * num_rows]
-
- # --- Process: evaluate each UDF row-by-row ---
- output_arrays = []
- for udf_func, offsets, zero_arg, arrow_return_type,
result_conv in udf_infos:
- rows = (
- [() for _ in range(num_rows)]
- if zero_arg
- else list(zip(*[columns[o] for o in offsets]))
- )
- results = _evaluate_batch_udf(udf_func, rows)
- verify_result_row_count(len(results), num_rows)
-
- # --- Output: Python -> Arrow ---
- converted = (
- [result_conv(r) for r in results] if result_conv is
not None else results
- )
- try:
- arr = pa.array(converted, type=arrow_return_type)
- except pa.lib.ArrowInvalid:
- arr = pa.array(converted).cast(
- target_type=arrow_return_type,
safe=runner_conf.safecheck
+ # Keep phase totals local so each metric is updated once instead
of per batch.
+ input_duration_ns = 0
+ udf_duration_ns = 0
+ output_duration_ns = 0
+ clock = time.perf_counter_ns
+ try:
+ for input_batch in data:
+ input_start_ns = clock()
+ # --- Input: Arrow -> Python columns ---
+ metrics.increment("pythonNumTimedBatches")
+ num_rows = input_batch.num_rows
+ columns = [
+ (
+ [conv(v) for v in
ArrowTableToRowsConversion._to_pylist(col)]
+ if conv is not None
+ else ArrowTableToRowsConversion._to_pylist(col)
)
- output_arrays.append(arr)
+ for col, conv in zip(input_batch.itercolumns(),
arrow_to_py_converters)
+ ]
+ if not columns:
+ columns = [[_NoValue] * num_rows]
+ output_arrays: list[pa.Array] = []
+ # The UDF loop replaces this when output preparation
begins.
+ output_start_ns = input_start_ns
+
+ # --- Process: evaluate each UDF row-by-row ---
+ for udf_func, offsets, zero_arg, arrow_return_type,
result_conv in udf_infos:
+ if output_arrays:
+ # One boundary ends the previous output and starts
the next input.
+ input_start_ns = clock()
+ output_duration_ns += input_start_ns -
output_start_ns
Review Comment:
Yeah it is a bit hard to follow. It reused one clock reading to end the
previous UDF’s output and start the next UDF’s input. I do this to save a clock
read, and the tradeoff is that it includes the loop transition to output time.
I removed the placeholder output_start_ns = input_start_ns and now record each
UDF’s output time in its own iteration. With one UDF, the worker still makes
four clock reads per batch. With multiple UDFs, it uses one additional read to
leave the loop time out. The final UDF’s output includes batch assembly. Do you
think this is better?
--
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]
