dongjoon-hyun opened a new pull request, #59597:
URL: https://github.com/apache/spark/pull/59597
### What changes were proposed in this pull request?
This PR fixes the `pythonDataSent` metric (`data sent to Python workers`) in
the pipelined Python UDF mode (`spark.python.udf.pipelined.enabled=true`),
which stops counting after a task sends about 2 GiB to its Python worker.
`PipelinedWriterRunnable` now creates a new `DataOutputStream` after every
buffer it flushes to the Python worker, as `ReaderInputStream` of the default
path does for every buffer refill.
### Why are the changes needed?
The runners measure `pythonDataSent` with `DataOutputStream.size()` deltas,
e.g., `PythonUDFRunner`, `BasicPythonArrowInput`, `BatchedPythonArrowInput`,
`GroupedPythonArrowInput`, `ColumnarArrowPythonInput`, and
`CoGroupedArrowPythonRunner`. `size()` is an `Int` that saturates at
`Int.MaxValue`. The default path passes a new `DataOutputStream` for every
buffer refill (64 KiB by default), so `size()` stays far below `Int.MaxValue`.
However, the pipelined mode (SPARK-56642) used one `DataOutputStream` for the
whole task. Once a task wrote about 2 GiB to it, the batch that crossed
`Int.MaxValue` was undercounted and every later batch added 0.
For example, `mapInArrow` over 2,500,000 rows of 1,000-byte strings in one
task reports the following `pythonDataSent` values.
| Mode | Before | After |
| - | -: | -: |
| Default | 2,510,678,000 | 2,510,678,000 |
| Pipelined | 2,147,477,526 | 2,510,678,000 |
Writers that keep an earlier stream, like the `ArrowStreamWriter` of the
stateful runners, still write to the same buffer. Those runners measure with
`ArrowStreamWriter.bytesWritten()` since SPARK-60313, so their metric does not
depend on the stream either.
### Does this PR introduce _any_ user-facing change?
Yes, only when `spark.python.udf.pipelined.enabled` is `true` (default:
`false`). For a task that sends more than 2 GiB to its Python worker, the `data
sent to Python workers` metric counts all the data instead of stopping at about
2 GiB. The pipelined mode is new in the unreleased Apache Spark 4.3.0
(SPARK-56642).
### How was this patch tested?
Added a test to `BasePythonRunnerSuite`. It runs `PipelinedWriterRunnable`
over a loopback socket with a writer that writes 300 bytes per call into
1,024-byte buffers, and checks that `DataOutputStream.size()` restarts for
every buffer and that all the data and `END_OF_STREAM` reach the socket.
Without this fix, it fails with `List(0, 300, 600, 900, 1200, ..., 5700) did
not equal List(0, 300, 600, 900, 0, 300, ...)`.
Also ran the following, and they pass.
```
build/sbt 'core/testOnly *BasePythonRunnerSuite'
TZ=UTC python/run-tests --testnames
'pyspark.sql.tests.pandas.test_pipelined_udf,pyspark.sql.tests.pandas.test_pandas_sqlmetrics'
```
The numbers above come from the following query with
`spark.python.udf.pipelined.enabled` set to `false` and `true` on `local[1]`.
```python
import pyarrow as pa
from pyspark.sql.functions import lit, repeat
def count_rows(batches):
for batch in batches:
yield pa.RecordBatch.from_pydict({"n": [batch.num_rows]})
df = (
spark.range(0, 2500000, 1, 1)
.select(repeat(lit("x"), 1000).alias("s"))
.mapInArrow(count_rows, "n long")
)
df.collect()
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
--
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]
