[SPARK-50642][CONNECT][SS]Fix the state schema for FlatMapGroupsWithState in spark connect when there is no initial state

[huanliwang-db]

In spark connect, when there is no initial state, we derived the state schema from the input:

create the initialDs from the original input: https://github.com/apache/spark/blob/master/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala#L679-L689
derived the state expression encoder from this initialDs which is incorrect: https://github.com/apache/spark/blob/master/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala#L729

Our unit tests fail to cover this case because it doesn't do the state update: https://github.com/apache/spark/blob/master/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsWithStateStreamingSuite.scala#L55-L59

after changing the stateFunc to the following

val stateFunc =
  (key: String, values: Iterator[ClickEvent], state: GroupState[ClickState]) => {
    if (state.exists) throw new IllegalArgumentException("state.exists should be false")
    val newState = ClickState(key, values.size)
    state.update(newState)
    Iterator(newState)
  }

the test is actually failing with

Cause: org.apache.spark.SparkException: Job aborted due to stage failure: Task 122 in stage 2.0 failed 1 times, most recent failure: Lost task 122.0 in stage 2.0 (TID 12) (192.168.68.84 executor driver): java.lang.ClassCastException: class org.apache.spark.sql.streaming.ClickState cannot be cast to class org.apache.spark.sql.streaming.ClickEvent (org.apache.spark.sql.streaming.ClickState and org.apache.spark.sql.streaming.ClickEvent are in unnamed module of loader 'app')
        at org.apache.spark.sql.catalyst.expressions.GeneratedClass$SpecificUnsafeProjection.CreateNamedStruct_0$(Unknown Source)

What changes were proposed in this pull request?

• introduce a new state_schema proto field
• pass the state agnostic encoder to the serialized udf
• pass the state schema to query proto for spark connect
• rebuild the state expression encoder based on the state agnostic encoder and state schema.

Why are the changes needed?

fix the broken behavior for flatMapGroupsWithState on spark connect

Does this PR introduce any user-facing change?

yes

How was this patch tested?

modified the existing unit tests.

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

no

[haiyangsun-db]

A better way to pass a schema (StructType) in spark connect than using json, see:

https://github.com/apache/spark/blob/master/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala#L197

.setStateSchema(DataTypeProtoConverter.toConnectProtoType(stateEncoder.schema))