fix

python/pyspark/sql/streaming/state_api_client.py

@@ -115,6 +115,8 @@ def _get_proto_state(self,
 
     def _send_proto_message(self, message: stateMessage.StateRequest) -> None:
         serialized_msg = message.SerializeToString()
+        # Writing zero here to indicate message version. This allows us to evolve the message
+        # format or even changing the message protocol in the future.
         write_int(0, self.sockfile)
         write_int(len(serialized_msg), self.sockfile)
         self.sockfile.write(serialized_msg)

python/pyspark/util.py

@@ -60,7 +60,7 @@
         PandasGroupedMapUDFWithStateType,
         ArrowGroupedMapUDFType,
         ArrowCogroupedMapUDFType,
-        PandasGroupedMapUDFTransformWithStateType,
+        PandasTransformWithStateUDFType,
     )
     from pyspark.sql._typing import (
         SQLArrowBatchedUDFType,
@@ -586,7 +586,7 @@ class PythonEvalType:
     SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE: "PandasGroupedMapUDFWithStateType" = 208
     SQL_GROUPED_MAP_ARROW_UDF: "ArrowGroupedMapUDFType" = 209
     SQL_COGROUPED_MAP_ARROW_UDF: "ArrowCogroupedMapUDFType" = 210
-    SQL_TRANSFORM_WITH_STATE: "PandasGroupedMapUDFTransformWithStateType" = 211
+    SQL_TRANSFORM_WITH_STATE_PANDAS_UDF: "PandasTransformWithStateUDFType" = 211
 
     SQL_TABLE_UDF: "SQLTableUDFType" = 300
     SQL_ARROW_TABLE_UDF: "SQLArrowTableUDFType" = 301

sql/core/src/main/scala/org/apache/spark/sql/RelationalGroupedDataset.scala

@@ -733,7 +733,6 @@ class RelationalGroupedDataset protected[sql](
    * along with corresponding state, and the results for all groups are combined into a new
    * [[DataFrame]].
    *
-   *
    * This function uses Apache Arrow as serialization format between Java executors and Python
    * workers.
    */

sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala

@@ -71,17 +71,17 @@ private[python] trait PythonArrowInput[IN] { self: BasePythonRunner[IN, _] =>
   protected val root = VectorSchemaRoot.create(arrowSchema, allocator)
   protected var writer: ArrowStreamWriter = _
 
-  protected def close(): Unit = {
-    Utils.tryWithSafeFinally {
-      // end writes footer to the output stream and doesn't clean any resources.
-      // It could throw exception if the output stream is closed, so it should be
-      // in the try block.
-      writer.end()
-    } {
-      root.close()
-      allocator.close()
-    }
+protected def close(): Unit = {
+  Utils.tryWithSafeFinally {
+    // end writes footer to the output stream and doesn't clean any resources.
+    // It could throw exception if the output stream is closed, so it should be
+    // in the try block.
+    writer.end()
+  } {
+    root.close()
+    allocator.close()
   }
+}
 
   protected override def newWriter(
       env: SparkEnv,

sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasExec.scala

@@ -75,7 +75,6 @@ case class TransformWithStateInPandasExec(
   private val groupingKeyExprEncoder = ExpressionEncoder(groupingKeySchema)
     .resolveAndBind().asInstanceOf[ExpressionEncoder[Any]]
 
-
   /** The keys that may have a watermark attribute. */
   override def keyExpressions: Seq[Attribute] = groupingAttributes
 

sql/core/src/main/scala/org/apache/spark/sql/execution/python/TransformWithStateInPandasStateServer.scala

@@ -34,24 +34,23 @@ import org.apache.spark.sql.streaming.ValueState
 import org.apache.spark.sql.types.StructType
 
 /**
- * This class is used to handle the state requests from the Python side.
+ * This class is used to handle the state requests from the Python side. It processes following
+ * state requests and return responses to the Python side:
+ * - Requests for managing explicit grouping key.
+ * - Stateful processor requests.
+ * - Requests for managing state variables (e.g. valueState).
  */
 class TransformWithStateInPandasStateServer(
     private val stateServerSocket: ServerSocket,
     private val statefulProcessorHandle: StatefulProcessorHandleImpl,
     private val groupingKeySchema: StructType)
-  extends Runnable
-  with Logging{
-
+  extends Runnable with Logging {
   private var inputStream: DataInputStream = _
   private var outputStream: DataOutputStream = _
-
   private val valueStates = mutable.HashMap[String, ValueState[Row]]()
 
   def run(): Unit = {
     val listeningSocket = stateServerSocket.accept()
-    logWarning(s"listening on socket - ${listeningSocket.getLocalAddress}")
-
     inputStream = new DataInputStream(
       new BufferedInputStream(listeningSocket.getInputStream))
     outputStream = new DataOutputStream(
@@ -60,21 +59,15 @@ class TransformWithStateInPandasStateServer(
 
     while (listeningSocket.isConnected &&
       statefulProcessorHandle.getHandleState != StatefulProcessorHandleState.CLOSED) {
-
       try {
         val version = inputStream.readInt()
-
         if (version != -1) {
           assert(version == 0)
           val messageLen = inputStream.readInt()
-
           val messageBytes = new Array[Byte](messageLen)
           inputStream.read(messageBytes)
-
           val message = StateRequest.parseFrom(ByteString.copyFrom(messageBytes))
-
           handleRequest(message)
-
           outputStream.flush()
         }
       } catch {