[SPARK-41431][CORE][SQL][UI] Protobuf serializer for SQLExecutionUIData

core/src/main/protobuf/org/apache/spark/status/protobuf/store_types.proto

@@ -258,3 +258,24 @@ message SpeculationStageSummaryWrapper {
   int32 stage_attempt_id = 2;
   SpeculationStageSummary info = 3;
 }
+
+message SQLPlanMetric {
+  string name = 1;
+  int64 accumulator_id = 2;
+  string metric_type = 3;
+}
+
+message SQLExecutionUIData {
+  int64 execution_id = 1;
+  string description = 2;
+  string details = 3;
+  string physical_plan_description = 4;
+  map<string, string> modified_configs = 5;
+  repeated SQLPlanMetric metrics = 6;
+  int64 submission_time = 7;
+  optional int64 completion_time = 8;
+  optional string error_message = 9;
+  map<int64, JobExecutionStatus> jobs = 10;
+  repeated int64 stages = 11;
+  map<int64, string> metric_values = 12;
+}

sql/core/pom.xml

@@ -147,6 +147,11 @@
       <groupId>org.apache.xbean</groupId>
       <artifactId>xbean-asm9-shaded</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.google.protobuf</groupId>
+      <artifactId>protobuf-java</artifactId>
+      <version>${protobuf.version}</version>
+    </dependency>
     <dependency>
       <groupId>org.scalacheck</groupId>
       <artifactId>scalacheck_${scala.binary.version}</artifactId>

sql/core/src/main/resources/META-INF/services/org.apache.spark.status.protobuf.ProtobufSerDe

@@ -0,0 +1,18 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+org.apache.spark.status.protobuf.sql.SQLExecutionUIDataSerializer

sql/core/src/main/scala/org/apache/spark/status/protobuf/sql/SQLExecutionUIDataSerializer.scala

@@ -0,0 +1,87 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.status.protobuf.sql
+
+import java.util.Date
+
+import collection.JavaConverters._
+
+import org.apache.spark.JobExecutionStatus
+import org.apache.spark.sql.execution.ui.SQLExecutionUIData
+import org.apache.spark.status.protobuf.{ProtobufSerDe, StoreTypes}
+import org.apache.spark.status.protobuf.Utils.getOptional
+
+class SQLExecutionUIDataSerializer extends ProtobufSerDe {
+
+  override val supportClass: Class[_] = classOf[SQLExecutionUIData]
+
+  override def serialize(input: Any): Array[Byte] = {
+    val ui = input.asInstanceOf[SQLExecutionUIData]
+    val builder = StoreTypes.SQLExecutionUIData.newBuilder()
+    builder.setExecutionId(ui.executionId)
+    builder.setDescription(ui.description)
+    builder.setDetails(ui.details)
+    builder.setPhysicalPlanDescription(ui.physicalPlanDescription)
+    ui.modifiedConfigs.foreach {
+      case (k, v) => builder.putModifiedConfigs(k, v)
+    }
+    ui.metrics.foreach(m => builder.addMetrics(SQLPlanMetricSerializer.serialize(m)))
+    builder.setSubmissionTime(ui.submissionTime)
+    ui.completionTime.foreach(ct => builder.setCompletionTime(ct.getTime))
+    ui.errorMessage.foreach(builder.setErrorMessage)
+    ui.jobs.foreach {
+      case (id, status) =>
+        builder.putJobs(id.toLong, StoreTypes.JobExecutionStatus.valueOf(status.toString))
+    }
+    ui.stages.foreach(stageId => builder.addStages(stageId.toLong))
+    ui.metricValues.foreach {
+      case (k, v) => builder.putMetricValues(k, v)
+    }
+    builder.build().toByteArray
+  }
+
+  override def deserialize(bytes: Array[Byte]): SQLExecutionUIData = {
+    val ui = StoreTypes.SQLExecutionUIData.parseFrom(bytes)
+    val completionTime =
+      getOptional(ui.hasCompletionTime, () => new Date(ui.getCompletionTime))
+    val errorMessage = getOptional(ui.hasErrorMessage, () => ui.getErrorMessage)
+    val metrics =
+      ui.getMetricsList.asScala.map(m => SQLPlanMetricSerializer.deserialize(m)).toSeq
+    val jobs = ui.getJobsMap.asScala.map {
+      case (jobId, status) => jobId.toInt -> JobExecutionStatus.valueOf(status.toString)
+    }.toMap
+    val metricValues = ui.getMetricValuesMap.asScala.map {
+      case (k, v) => k.toLong -> v
+    }.toMap
+
+    new SQLExecutionUIData(
+      executionId = ui.getExecutionId,
+      description = ui.getDescription,
+      details = ui.getDetails,
+      physicalPlanDescription = ui.getPhysicalPlanDescription,
+      modifiedConfigs = ui.getModifiedConfigsMap.asScala.toMap,
+      metrics = metrics,
+      submissionTime = ui.getSubmissionTime,
+      completionTime = completionTime,
+      errorMessage = errorMessage,
+      jobs = jobs,
+      stages = ui.getStagesList.asScala.map(_.toInt).toSet,
+      metricValues = metricValues
+    )
+  }
+}

sql/core/src/main/scala/org/apache/spark/status/protobuf/sql/SQLPlanMetricSerializer.scala

@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.status.protobuf.sql
+
+import org.apache.spark.sql.execution.ui.SQLPlanMetric
+import org.apache.spark.status.protobuf.StoreTypes
+
+object SQLPlanMetricSerializer {
+
+  def serialize(metric: SQLPlanMetric): StoreTypes.SQLPlanMetric = {
+    StoreTypes.SQLPlanMetric.newBuilder()
+      .setName(metric.name)
+      .setAccumulatorId(metric.accumulatorId)
+      .setMetricType(metric.metricType)
+      .build()
+  }
+
+  def deserialize(metrics: StoreTypes.SQLPlanMetric): SQLPlanMetric = {
+    SQLPlanMetric(metrics.getName, metrics.getAccumulatorId, metrics.getMetricType)
+  }
+}

sql/core/src/test/scala/org/apache/spark/status/protobuf/sql/KVStoreProtobufSerializerSuite.scala

@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.status.protobuf.sql
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.execution.ui.SQLExecutionUIData
+import org.apache.spark.status.api.v1.sql.SqlResourceSuite
+import org.apache.spark.status.protobuf.KVStoreProtobufSerializer
+
+class KVStoreProtobufSerializerSuite extends SparkFunSuite {
+
+  private val serializer = new KVStoreProtobufSerializer()
+
+  test("SQLExecutionUIData") {
+    val input = SqlResourceSuite.sqlExecutionUIData
+    val bytes = serializer.serialize(input)
+    val result = serializer.deserialize(bytes, classOf[SQLExecutionUIData])
+    assert(result.executionId == input.executionId)
+    assert(result.description == input.description)
+    assert(result.details == input.details)
+    assert(result.physicalPlanDescription == input.physicalPlanDescription)
+    assert(result.modifiedConfigs == input.modifiedConfigs)
+    assert(result.metrics == input.metrics)
+    assert(result.submissionTime == input.submissionTime)
+    assert(result.completionTime == input.completionTime)
+    assert(result.errorMessage == input.errorMessage)
+    assert(result.jobs == input.jobs)
+    assert(result.stages == input.stages)
+    assert(result.metricValues == input.metricValues)
+  }
+}