dongjoon-hyun opened a new pull request, #59570: URL: https://github.com/apache/spark/pull/59570
### What changes were proposed in this pull request? This PR makes the following runners measure the `pythonDataSent` metric with `ArrowStreamWriter.bytesWritten()` around each whole `writeNextBatchToArrowStream` call, instead of `dataOut.size()` around part of it. - `ApplyInPandasWithStatePythonRunner` (`applyInPandasWithState`) - `TransformWithStateInPySparkPythonRunner` (`transformWithState` and `transformWithStateInPandas`) - `TransformWithStateInPySparkPythonInitialStateRunner` (the same with an initial state) ### Why are the changes needed? These runners write all their Arrow batches through one `ArrowStreamWriter`, which `PythonArrowInput` creates over the `dataOut` of the first call. Since SPARK-44705, the default (non-pipelined) path passes a new `DataOutputStream` on every buffer refill (64 KiB by default). The batches still reach the same buffer through the first stream, but `dataOut.size()` of the later streams does not change. So the metric counts only the first refill. Two calls were also never measured: the last call of `applyInPandasWithState`, which flushes its last batch, and the call of the initial state runner that flushes its last initial state batch. So `applyInPandasWithState` reports 0 B for a task whose input fits in one batch. ### Does this PR introduce _any_ user-facing change? Yes. The `data sent to Python workers` metric of `FlatMapGroupsInPandasWithState` and `TransformWithStateInPySpark` now counts all the data. For example, in a manual end-to-end check, an `applyInPandasWithState` micro-batch of 1,000 rows reported 0 B instead of 134,424 B, and one of 100,000 rows reported 1,313,448 B instead of 13,134,488 B. The undercount has existed since: - Apache Spark [v3.4.0](https://github.com/apache/spark/releases/tag/v3.4.0) (2023-04-07) for the last batch of `applyInPandasWithState` (SPARK-34265) - Apache Spark [v4.0.0](https://github.com/apache/spark/releases/tag/v4.0.0) (2025-05-19) for the buffer refills (SPARK-44705) - Apache Spark [v4.2.0](https://github.com/apache/spark/releases/tag/v4.2.0) (2026-07-11) for the last initial state batch (SPARK-54392) ### How was this patch tested? Added `StatefulPythonRunnerDataSentSuite`. It drives each runner's input writer without a Python worker, with a new `DataOutputStream` per call and with one shared stream (pipelined mode), and checks that `pythonDataSent` equals the bytes written after the Arrow schema. All three tests fail without this fix. Also ran the following, and they pass. ``` build/sbt 'sql/testOnly *StatefulPythonRunnerDataSentSuite *BaseStreamingArrowWriterSuite *FlatMapGroupsInPandasWithStateSuite *FlatMapGroupsInPandasWithStateDistributionSuite' TZ=UTC python/run-tests --testnames 'pyspark.sql.tests.pandas.streaming.test_pandas_transform_with_state,pyspark.sql.tests.pandas.streaming.test_transform_with_state,pyspark.sql.tests.pandas.test_pandas_grouped_map_with_state' ``` ### 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]
