[SPARK-46124][FOLLOWUP][CONNECT][SS] Send missing fields in StreamingQueryProgress to client

[WweiL]

What changes were proposed in this pull request?

Currently in PySpark Client, calling query.lastProgress won't return you three fields:

spark/sql/core/src/main/scala/org/apache/spark/sql/streaming/progress.scala

Lines 145 to 152 in 0bc2c69

	/** The aggregate (across all sources) number of records processed in a trigger. */
	def numInputRows: Long = sources.map(_.numInputRows).sum
	/** The aggregate (across all sources) rate of data arriving. */
	def inputRowsPerSecond: Double = sources.map(_.inputRowsPerSecond).sum
	/** The aggregate (across all sources) rate at which Spark is processing data. */
	def processedRowsPerSecond: Double = sources.map(_.processedRowsPerSecond).sum

This is because these three fields aren't native fields and are supposed to be reconstructed when rebuilding the client side StreamingQueryProgress. But this brings a bug only in Python Spark Connect.

This is not an issue for scala connect client, because the client side reconstruct the StreamingQueryProgress using the same way it's constructed in the server side.

The problem with python is the return type is a dict, so there is no such reconstruction, and therefore these fields are not there.

spark/python/pyspark/sql/connect/streaming/query.py

Line 122 in d6fcba0

def lastProgress(self) -> Optional[Dict[str, Any]]:

This is not an issue in classic pyspark also, because classic pyspark calls the jsonValue of StreamingQueryProgress, which has all fields defined:

spark/sql/core/src/main/scala/org/apache/spark/sql/streaming/progress.scala

Lines 162 to 178 in 0bc2c69

	private[sql] def jsonValue: JValue = {
	("id" -> JString(id.toString)) ~
	("runId" -> JString(runId.toString)) ~
	("name" -> JString(name)) ~
	("timestamp" -> JString(timestamp)) ~
	("batchId" -> JInt(batchId)) ~
	("batchDuration" -> JInt(batchDuration)) ~
	("numInputRows" -> JInt(numInputRows)) ~
	("inputRowsPerSecond" -> safeDoubleToJValue(inputRowsPerSecond)) ~
	("processedRowsPerSecond" -> safeDoubleToJValue(processedRowsPerSecond)) ~
	("durationMs" -> safeMapToJValue[JLong](durationMs, v => JInt(v.toLong))) ~
	("eventTime" -> safeMapToJValue[String](eventTime, s => JString(s))) ~
	("stateOperators" -> JArray(stateOperators.map(_.jsonValue).toList)) ~
	("sources" -> JArray(sources.map(_.jsonValue).toList)) ~
	("sink" -> sink.jsonValue) ~
	("observedMetrics" -> safeMapToJValue[Row](observedMetrics, row => row.jsonValue))
	}

In this PR we send all such fields to the client to fix this issue. This won't cause backward compatibility issue because that field is a map anyways.

Why are the changes needed?

Bug fix

Does this PR introduce any user-facing change?

No

How was this patch tested?

Added unit test, passed locally. Before the fix, the test would fail with:

======================================================================
ERROR: test_streaming_last_progress (pyspark.sql.tests.connect.streaming.test_parity_streaming.StreamingParityTests.test_streaming_last_progress)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/wei.liu/oss-spark/python/pyspark/sql/tests/streaming/test_streaming.py", line 68, in test_streaming_last_progress
    progress = StreamingQueryProgress.fromJson(query.lastProgress)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/wei.liu/oss-spark/python/pyspark/sql/streaming/listener.py", line 489, in fromJson
    numInputRows=j["numInputRows"],
                 ~^^^^^^^^^^^^^^^^
KeyError: 'numInputRows'

Was this patch authored or co-authored using generative AI tooling?

No

[WweiL]

cc @HeartSaVioR @HyukjinKwon @LuciferYang, can you guys take a look? Thank you!

[WweiL]

this is just a minor fix it will make the test more stable. We call query.processAllAvailable below, which has a very small chance of being a indefinite blocking call on rate source.

[LuciferYang]

In the original design, methods jsonString and fromJson are paired, so the test failure in StreamingQueryStatusAndProgressSuite should be related to the current pr:

[LuciferYang]

Moreover, the test failures in ClientStreamingQuerySuite may also be related to the current PR, seems StreamingQueryProgress.observedMetrics is null with this pr(it's a empty Map without this pr)

spark/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/streaming/ClientStreamingQuerySuite.scala

Line 108 in ab00533

assert(lastProgress.observedMetrics.isEmpty)

Does this mean that the new changes are not fully compatible with previous behaviors? @WweiL @HyukjinKwon

[WweiL]

Ah thanks, I only tested the new test case, glad existing tests catch them, let me fix! @LuciferYang

[WweiL]

Now I actually think we should fix python client side reconstruction

[WweiL]

That three fields are just constructed using other existing info:

spark/sql/core/src/main/scala/org/apache/spark/sql/streaming/progress.scala

Lines 145 to 152 in 0bc2c69

	/** The aggregate (across all sources) number of records processed in a trigger. */
	def numInputRows: Long = sources.map(_.numInputRows).sum
	/** The aggregate (across all sources) rate of data arriving. */
	def inputRowsPerSecond: Double = sources.map(_.inputRowsPerSecond).sum
	/** The aggregate (across all sources) rate at which Spark is processing data. */
	def processedRowsPerSecond: Double = sources.map(_.processedRowsPerSecond).sum

There is no absolute reason we also send them from the server.

I can just force the client to reconstruct these fields as in the scala definition.

But there is always this API difference, in python, the return type of lastProgress is Dict, but in scala it's the actual StreamingQueryProgress object:

spark/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQuery.scala

Lines 94 to 101 in 1a5d22a

	def recentProgress: Array[StreamingQueryProgress]
	/**
	* Returns the most recent [[StreamingQueryProgress]] update of this streaming query.
	*
	* @since 2.1.0
	*/
	def lastProgress: StreamingQueryProgress

I feel that we could let the python client also return that 'StreamingQueryProgress' object, and put the reconstruct logic inside the client side StreamingQueryProgress, what do you guys think?
@LuciferYang @HyukjinKwon @HeartSaVioR

[HyukjinKwon]

We should have done that. I wonder if we can have a class that implements dict e.g., StreamingQueryProgress(dict) to avoid a breaking change.

[WweiL]

Created another PR: #46921, it is not ready for review yet because I found another bug when working on it. After #46920 is merged, it will be ready for review