gaogaotiantian commented on code in PR #59246:
URL: https://github.com/apache/spark/pull/59246#discussion_r4236167826
##########
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 ---
Review Comment:
This comment should either be before `clock()` or after `increment`. It's
weird to have it in the middle.
##########
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:
I don't quite understand what are we trying to do here. What's this value?
`output_start_ns` is from the previous iteration of udf func? This is super
difficult to follow. Why do we need this? Doesn't this include some random time
outside of this function?
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -213,21 +213,12 @@ private[spark] object BasePythonRunner extends Logging {
}
}
- // Values used to update the existing Python SQL metrics and task spill
metrics.
- private[python] case class WorkerMetrics(
- bootTimestampMs: Long,
- initTimestampMs: Long,
- finishTimestampMs: Long,
- pythonExecutionDurationMs: Long,
- memoryBytesSpilled: Long,
- diskBytesSpilled: Long)
-
private lazy val workerMetricsMapper = new ObjectMapper()
.enable(JsonParser.Feature.STRICT_DUPLICATE_DETECTION)
.enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS)
- /** Read and validate worker metrics after METRICS_DATA; ignore additional
fields. */
- private[python] def readWorkerMetrics(stream: DataInputStream):
WorkerMetrics = {
+ /** Decode all numeric fields after METRICS_DATA without knowing their
metric names. */
+ private[python] def readWorkerMetrics(stream: DataInputStream): Map[String,
Long] = {
Review Comment:
Is there a restriction here for all the values to be the same type? Is it
possible for us to emit string in the future?
--
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]
