[SPARK-41432][UI][SQL] Protobuf serializer for SparkPlanGraphWrapper

project/SparkBuild.scala

@@ -437,7 +437,9 @@ object SparkBuild extends PomBuild {
   /* Catalyst ANTLR generation settings */
   enable(Catalyst.settings)(catalyst)
 
-  /* Spark SQL Core console settings */
+  /* Spark SQL Core console settings
+   * Generate and pick the spark build info from extra-resources
+   */
   enable(SQL.settings)(sql)
 
   /* Hive console settings */
@@ -825,6 +827,7 @@ object SparkConnect {
   }
 }
 
+
 object SparkProtobuf {
   import BuildCommons.protoVersion
 
@@ -1069,6 +1072,9 @@ object Catalyst {
 }
 
 object SQL {
+
+  import BuildCommons.protoVersion
+
   lazy val settings = Seq(
     (console / initialCommands) :=
       """
@@ -1091,7 +1097,30 @@ object SQL {
         |import sqlContext._
       """.stripMargin,
     (console / cleanupCommands) := "sc.stop()"
-  )
+  ) ++ Seq(
+    // Setting version for the protobuf compiler. This has to be propagated to every sub-project
+    // even if the project is not using it.
+    PB.protocVersion := BuildCommons.protoVersion,
+    // For some reason the resolution from the imported Maven build does not work for some
+    // of these dependendencies that we need to shade later on.
+    libraryDependencies ++= {
+      Seq(
+        "com.google.protobuf" % "protobuf-java" % protoVersion % "protobuf"
+      )
+    },
+    (Compile / PB.targets) := Seq(
+      PB.gens.java -> (Compile / sourceManaged).value
+    )
+  ) ++ {
+    val sparkProtocExecPath = sys.props.get("spark.protoc.executable.path")
+    if (sparkProtocExecPath.isDefined) {
+      Seq(
+        PB.protocExecutable := file(sparkProtocExecPath.get)
+      )
+    } else {
+      Seq.empty
+    }
+  }
 }
 
 object Hive {

sql/core/pom.xml

@@ -208,6 +208,12 @@
       <artifactId>htmlunit-driver</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.google.protobuf</groupId>
+      <artifactId>protobuf-java</artifactId>
+      <version>${protobuf.version}</version>
+      <scope>compile</scope>
+    </dependency>
   </dependencies>
   <build>
     <outputDirectory>target/scala-${scala.binary.version}/classes</outputDirectory>
@@ -272,6 +278,27 @@
           </execution>
         </executions>
       </plugin>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-shade-plugin</artifactId>
+        <configuration>
+          <shadedArtifactAttached>false</shadedArtifactAttached>
+          <shadeTestJar>true</shadeTestJar>
+          <relocations>
+            <relocation>
+              <pattern>com.google.common</pattern>
+              <shadedPattern>${spark.shade.packageName}.guava</shadedPattern>
+            </relocation>
+            <relocation>
+              <pattern>com.google.protobuf</pattern>
+              <shadedPattern>${spark.shade.packageName}.spark-core.protobuf</shadedPattern>
+              <includes>
+                <include>com.google.protobuf.**</include>
+              </includes>
+            </relocation>
+          </relocations>
+        </configuration>
+      </plugin>
     </plugins>
   </build>
 
@@ -291,5 +318,68 @@
         </dependency>
       </dependencies>
     </profile>
+    <profile>
+      <id>default-protoc</id>
+      <activation>
+        <property>
+          <name>!skipDefaultProtoc</name>
+        </property>
+      </activation>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>com.github.os72</groupId>
+            <artifactId>protoc-jar-maven-plugin</artifactId>
+            <version>${protoc-jar-maven-plugin.version}</version>
+            <executions>
+              <execution>
+                <phase>generate-sources</phase>
+                <goals>
+                  <goal>run</goal>
+                </goals>
+                <configuration>
+                  <protocArtifact>com.google.protobuf:protoc:${protobuf.version}</protocArtifact>
+                  <protocVersion>${protobuf.version}</protocVersion>
+                  <inputDirectories>
+                    <include>src/main/protobuf</include>
+                  </inputDirectories>
+                </configuration>
+              </execution>
+            </executions>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
+    <profile>
+      <id>user-defined-protoc</id>
+      <properties>
+        <spark.protoc.executable.path>${env.SPARK_PROTOC_EXEC_PATH}</spark.protoc.executable.path>
+      </properties>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>com.github.os72</groupId>
+            <artifactId>protoc-jar-maven-plugin</artifactId>
+            <version>${protoc-jar-maven-plugin.version}</version>
+            <executions>
+              <execution>
+                <phase>generate-sources</phase>
+                <goals>
+                  <goal>run</goal>
+                </goals>
+                <configuration>
+                  <protocArtifact>com.google.protobuf:protoc:${protobuf.version}</protocArtifact>
+                  <protocVersion>${protobuf.version}</protocVersion>
+                  <protocCommand>${spark.protoc.executable.path}</protocCommand>
+                  <inputDirectories>
+                    <include>src/main/protobuf</include>
+                  </inputDirectories>
+                </configuration>
+              </execution>
+            </executions>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
   </profiles>
 </project>

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

@@ -0,0 +1,56 @@
+/*
+ * 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.
+ */
+
+syntax = "proto3";
+package org.apache.spark.status.protobuf.sql;
+
+message SQLPlanMetric {
+  string name = 1;
+  int64 accumulator_id = 2;
+  string metric_type = 3;
+}
+
+message SparkPlanGraphNode {
+  int64 id = 1;
+  string name = 2;
+  string desc = 3;
+  repeated SQLPlanMetric metrics = 4;
+}
+
+message SparkPlanGraphClusterWrapper {
+  int64 id = 1;
+  string name = 2;
+  string desc = 3;
+  repeated SparkPlanGraphNodeWrapper nodes = 4;
+  repeated SQLPlanMetric metrics = 5;
+}
+
+message SparkPlanGraphNodeWrapper {
+  SparkPlanGraphNode node = 1;
+  SparkPlanGraphClusterWrapper cluster = 2;
+}
+
+message SparkPlanGraphEdge {
+  int64 from_id = 1;
+  int64 to_id = 2;
+}
+
+message SparkPlanGraphWrapper {
+  int64 execution_id = 1;
+  repeated SparkPlanGraphNodeWrapper nodes = 2;
+  repeated SparkPlanGraphEdge edges = 3;
+}

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.SparkPlanGraphWrapperSerializer

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

@@ -0,0 +1,151 @@
+/*
+ * 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 collection.JavaConverters._
+
+import org.apache.spark.sql.execution.ui.{SparkPlanGraphClusterWrapper, SparkPlanGraphEdge, SparkPlanGraphNode, SparkPlanGraphNodeWrapper, SparkPlanGraphWrapper, SQLPlanMetric}
+import org.apache.spark.status.protobuf.ProtobufSerDe
+
+class SparkPlanGraphWrapperSerializer extends ProtobufSerDe {
+
+  override val supportClass: Class[_] = classOf[SparkPlanGraphWrapper]
+
+  override def serialize(input: Any): Array[Byte] =
+    serialize(input.asInstanceOf[SparkPlanGraphWrapper])
+
+  private def serialize(plan: SparkPlanGraphWrapper): Array[Byte] = {
+    val builder = StoreTypes.SparkPlanGraphWrapper.newBuilder()
+    builder.setExecutionId(plan.executionId)
+    plan.nodes.foreach { node =>
+      builder.addNodes(serializeSparkPlanGraphNodeWrapper(node))
+    }
+    plan.edges.foreach {edge =>
+      builder.addEdges(serializeSparkPlanGraphEdge(edge))
+    }
+    builder.build().toByteArray
+  }
+
+  def deserialize(bytes: Array[Byte]): SparkPlanGraphWrapper = {
+    val wrapper = StoreTypes.SparkPlanGraphWrapper.parseFrom(bytes)
+    new SparkPlanGraphWrapper(
+      executionId = wrapper.getExecutionId,
+      nodes = wrapper.getNodesList.asScala.map(deserializeSparkPlanGraphNodeWrapper).toSeq,
+      edges = wrapper.getEdgesList.asScala.map(deserializeSparkPlanGraphEdge).toSeq
+    )
+  }
+
+  private def serializeSparkPlanGraphNodeWrapper(input: SparkPlanGraphNodeWrapper):
+    StoreTypes.SparkPlanGraphNodeWrapper = {
+
+    val builder = StoreTypes.SparkPlanGraphNodeWrapper.newBuilder()
+    builder.setNode(serializeSparkPlanGraphNode(input.node))
+    builder.setCluster(serializeSparkPlanGraphClusterWrapper(input.cluster))
+    builder.build()
+  }
+
+  private def deserializeSparkPlanGraphNodeWrapper(input: StoreTypes.SparkPlanGraphNodeWrapper):
+    SparkPlanGraphNodeWrapper = {
+
+    new SparkPlanGraphNodeWrapper(
+      node = deserializeSparkPlanGraphNode(input.getNode),
+      cluster = deserializeSparkPlanGraphClusterWrapper(input.getCluster)
+    )
+  }
+
+  private def serializeSparkPlanGraphEdge(edge: SparkPlanGraphEdge):
+    StoreTypes.SparkPlanGraphEdge = {
+    val builder = StoreTypes.SparkPlanGraphEdge.newBuilder()
+    builder.setFromId(edge.fromId)
+    builder.setToId(edge.toId)
+    builder.build()
+  }
+
+  private def deserializeSparkPlanGraphEdge(edge: StoreTypes.SparkPlanGraphEdge):
+    SparkPlanGraphEdge = {
+    SparkPlanGraphEdge(
+      fromId = edge.getFromId,
+      toId = edge.getToId)
+  }
+
+  private def serializeSparkPlanGraphNode(node: SparkPlanGraphNode):
+    StoreTypes.SparkPlanGraphNode = {
+    val builder = StoreTypes.SparkPlanGraphNode.newBuilder()
+    builder.setId(node.id)
+    builder.setName(node.name)
+    builder.setName(node.desc)
+    node.metrics.foreach { metric =>
+      builder.addMetrics(serializeSQLPlanMetric(metric))
+    }
+    builder.build()
+  }
+
+  private def deserializeSparkPlanGraphNode(node: StoreTypes.SparkPlanGraphNode):
+    SparkPlanGraphNode = {
+
+    new SparkPlanGraphNode(
+      id = node.getId,
+      name = node.getName,
+      desc = node.getDesc,
+      metrics = node.getMetricsList.asScala.map(deserializeSQLPlanMetric).toSeq
+    )
+  }
+
+  private def serializeSparkPlanGraphClusterWrapper(cluster: SparkPlanGraphClusterWrapper):
+    StoreTypes.SparkPlanGraphClusterWrapper = {
+    val builder = StoreTypes.SparkPlanGraphClusterWrapper.newBuilder()
+    builder.setId(cluster.id)
+    builder.setName(cluster.name)
+    builder.setDesc(cluster.desc)
+    cluster.nodes.foreach { node =>
+      builder.addNodes(serializeSparkPlanGraphNodeWrapper(node))
+    }
+    cluster.metrics.foreach { metric =>
+      builder.addMetrics(serializeSQLPlanMetric(metric))
+    }
+    builder.build()
+  }
+
+  private def deserializeSparkPlanGraphClusterWrapper(
+    cluster: StoreTypes.SparkPlanGraphClusterWrapper): SparkPlanGraphClusterWrapper = {
+
+    new SparkPlanGraphClusterWrapper(
+      id = cluster.getId,
+      name = cluster.getName,
+      desc = cluster.getDesc,
+      nodes = cluster.getNodesList.asScala.map(deserializeSparkPlanGraphNodeWrapper).toSeq,
+      metrics = cluster.getMetricsList.asScala.map(deserializeSQLPlanMetric).toSeq
+    )
+  }
+
+  private def serializeSQLPlanMetric(metric: SQLPlanMetric): StoreTypes.SQLPlanMetric = {
+    val builder = StoreTypes.SQLPlanMetric.newBuilder()
+    builder.setName(metric.name)
+    builder.setAccumulatorId(metric.accumulatorId)
+    builder.setMetricType(metric.metricType)
+    builder.build()
+  }
+
+  private def deserializeSQLPlanMetric(metric: StoreTypes.SQLPlanMetric): SQLPlanMetric = {
+    SQLPlanMetric(
+      name = metric.getName,
+      accumulatorId = metric.getAccumulatorId,
+      metricType = metric.getMetricType
+    )
+  }
+}