[KYUUBI #6577] Define a PlanOnlyExecutor interface to extend plan onl…

…y mode of spark engine using SPI
apache · Aug 2, 2024 · c3ce8dd · c3ce8dd
1 parent c94f0d7
commit c3ce8dd
Show file tree

Hide file tree

Showing 6 changed files with 131 additions and 0 deletions.
diff --git a/extensions/spark/kyuubi-spark-engine-plugin/pom.xml b/extensions/spark/kyuubi-spark-engine-plugin/pom.xml
@@ -0,0 +1,38 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0"
+         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+    <parent>
+        <groupId>org.apache.kyuubi</groupId>
+        <artifactId>kyuubi-parent</artifactId>
+        <version>1.10.0-SNAPSHOT</version>
+        <relativePath>../../../pom.xml</relativePath>
+    </parent>
+
+    <artifactId>kyuubi-spark-engine-plugin</artifactId>
+    <packaging>jar</packaging>
+    <name>Kyuubi Project Spark Engine Plugin</name>
+    <url>https://kyuubi.apache.org/</url>
+
+    <dependencies>
+        <dependency>
+            <groupId>org.apache.spark</groupId>
+            <artifactId>spark-sql_${scala.binary.version}</artifactId>
+            <scope>provided</scope>
+        </dependency>
+    </dependencies>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>net.alchim31.maven</groupId>
+                <artifactId>scala-maven-plugin</artifactId>
+                <configuration>
+                    <skip>true</skip>
+                </configuration>
+            </plugin>
+        </plugins>
+    </build>
+
+</project>
diff --git a/...k-engine-plugin/src/main/java/org/apache/kyuubi/engine/spark/plugin/PlanOnlyExecutor.java b/...k-engine-plugin/src/main/java/org/apache/kyuubi/engine/spark/plugin/PlanOnlyExecutor.java
@@ -0,0 +1,41 @@
+/*
+ * 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.kyuubi.engine.spark.plugin;
+
+import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan;
+
+public interface PlanOnlyExecutor {
+
+    String mode();
+
+    default String execute(SparkSession spark, LogicalPlan parsedPlan) {
+        return execute(spark, parsedPlan, "plain");
+    }
+
+    /**
+     * execute the parsed plan with plan only mode
+     *
+     * @param spark       the spark session
+     * @param parsedPlan  the parsed plan
+     * @param style       the style of the result (plain or json)
+     * @return the result in the specified style
+     */
+    String execute(SparkSession spark, LogicalPlan parsedPlan, String style);
+
+}
diff --git a/externals/kyuubi-spark-sql-engine/pom.xml b/externals/kyuubi-spark-sql-engine/pom.xml
@@ -49,6 +49,12 @@
             <version>${project.version}</version>
         </dependency>
 
+        <dependency>
+            <groupId>org.apache.kyuubi</groupId>
+            <artifactId>kyuubi-spark-engine-plugin</artifactId>
+            <version>${project.version}</version>
+        </dependency>
+
         <dependency>
             <groupId>io.grpc</groupId>
             <artifactId>grpc-core</artifactId>

diff --git a/...ql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala b/...ql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala
@@ -17,6 +17,8 @@
 
 package org.apache.kyuubi.engine.spark.operation
 
+import java.util.Locale
+
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.fasterxml.jackson.module.scala.DefaultScalaModule
 import org.apache.spark.kyuubi.SparkUtilsHelper
@@ -28,6 +30,7 @@ import org.apache.spark.sql.types.StructType
 import org.apache.kyuubi.KyuubiSQLException
 import org.apache.kyuubi.config.KyuubiConf.{LINEAGE_PARSER_PLUGIN_PROVIDER, OPERATION_PLAN_ONLY_EXCLUDES, OPERATION_PLAN_ONLY_OUT_STYLE}
 import org.apache.kyuubi.engine.spark.KyuubiSparkUtil.getSessionConf
+import org.apache.kyuubi.engine.spark.util.PlanOnlyExecutors
 import org.apache.kyuubi.operation.{AnalyzeMode, ArrayFetchIterator, ExecutionMode, IterableFetchIterator, JsonStyle, LineageMode, OperationHandle, OptimizeMode, OptimizeWithStatsMode, ParseMode, PhysicalMode, PlainStyle, PlanOnlyMode, PlanOnlyStyle, UnknownMode, UnknownStyle}
 import org.apache.kyuubi.operation.PlanOnlyMode.{notSupportedModeError, unknownModeError}
 import org.apache.kyuubi.operation.PlanOnlyStyle.{notSupportedStyleError, unknownStyleError}
@@ -150,6 +153,9 @@ class PlanOnlyStatement(
       case LineageMode =>
         val result = parseLineage(spark, plan)
         iter = new IterableFetchIterator(Seq(Row(result)))
+      case PlanOnlyExecutors(executor) =>
+        val result = executor.execute(spark, plan, style.name.toLowerCase(Locale.ROOT))
+        iter = new IterableFetchIterator(Seq(Row(result)))
       case UnknownMode => throw unknownModeError(mode)
       case _ =>
         throw KyuubiSQLException(s"The operation mode $mode" +

diff --git a/...ark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/util/PlanOnlyExecutors.scala b/...ark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/util/PlanOnlyExecutors.scala
@@ -0,0 +1,39 @@
+/*
+ * 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.kyuubi.engine.spark.util
+
+import java.util.Locale
+
+import org.apache.kyuubi.engine.spark.plugin.PlanOnlyExecutor
+import org.apache.kyuubi.operation.PlanOnlyMode
+import org.apache.kyuubi.util.reflect.ReflectUtils
+
+object PlanOnlyExecutors {
+
+  private lazy val executors: Map[String, PlanOnlyExecutor] = {
+    ReflectUtils.loadFromServiceLoader[PlanOnlyExecutor]()
+      .map(e => e.mode().toLowerCase(Locale.ROOT) -> e).toMap
+  }
+
+  def getExecutor(mode: PlanOnlyMode): Option[PlanOnlyExecutor] = {
+    executors.get(mode.name.toLowerCase(Locale.ROOT))
+  }
+
+  def unapply(mode: PlanOnlyMode): Option[PlanOnlyExecutor] = getExecutor(mode)
+
+}
diff --git a/pom.xml b/pom.xml
@@ -60,6 +60,7 @@
         <module>extensions/spark/kyuubi-spark-connector-common</module>
         <module>extensions/spark/kyuubi-spark-connector-tpcds</module>
         <module>extensions/spark/kyuubi-spark-connector-tpch</module>
+        <module>extensions/spark/kyuubi-spark-engine-plugin</module>
         <module>extensions/spark/kyuubi-spark-lineage</module>
         <module>externals/kyuubi-chat-engine</module>
         <module>externals/kyuubi-download</module>