From dc742bf8aa8a57b179052c8af1ccbc17cba1483a Mon Sep 17 00:00:00 2001
From: wforget <643348094@qq.com>
Date: Thu, 16 Nov 2023 14:37:17 +0800
Subject: [PATCH] Add ReplaceIcebergTableParser Rule

---
 .../spark/kyuubi-extension-spark-3-5/pom.xml  |   6 +
 .../org/apache/kyuubi/sql/KyuubiSQLConf.scala |   6 +
 .../sql/KyuubiSparkSQLCommonExtension.scala   |   3 +
 .../parser/ReplaceIcebergTableParser.scala    | 155 +++++++++++++
 .../spark/sql/ReplaceIcebergTableSuite.scala  | 215 ++++++++++++++++++
 5 files changed, 385 insertions(+)
 create mode 100644 extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/spark/sql/catalyst/parser/ReplaceIcebergTableParser.scala
 create mode 100644 extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/ReplaceIcebergTableSuite.scala

diff --git a/extensions/spark/kyuubi-extension-spark-3-5/pom.xml b/extensions/spark/kyuubi-extension-spark-3-5/pom.xml
index d55501ce860..0bf422881ab 100644
--- a/extensions/spark/kyuubi-extension-spark-3-5/pom.xml
+++ b/extensions/spark/kyuubi-extension-spark-3-5/pom.xml
@@ -105,6 +105,12 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>org.apache.iceberg</groupId>
+            <artifactId>iceberg-spark-runtime-${spark.binary.version}_${scala.binary.version}</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <!--
           Spark requires `commons-collections` and `commons-io` but got them from transitive
           dependencies of `hadoop-client`. As we are using Hadoop Shaded Client, we need add
diff --git a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala
index 5ac144de6e7..932d45cb9a1 100644
--- a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala
+++ b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSQLConf.scala
@@ -288,4 +288,10 @@ object KyuubiSQLConf {
       .booleanConf
       .createWithDefault(false)
 
+  val ICEBERG_TABLE_IMPLICIT_CONVERSION =
+    buildConf("spark.sql.qiyi.optimizer.icebergTableImplicitConversion")
+      .doc("Implicit conversion of iceberg table to automatically add `iceberg_catalog`.")
+      .version("1.5.0")
+      .booleanConf
+      .createWithDefault(false)
 }
diff --git a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLCommonExtension.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLCommonExtension.scala
index 22d253b6f4f..7238fb28f68 100644
--- a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLCommonExtension.scala
+++ b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/kyuubi/sql/KyuubiSparkSQLCommonExtension.scala
@@ -18,6 +18,7 @@
 package org.apache.kyuubi.sql
 
 import org.apache.spark.sql.SparkSessionExtensions
+import org.apache.spark.sql.catalyst.parser.ReplaceIcebergTableParser
 
 import org.apache.kyuubi.sql.zorder.{InsertZorderBeforeWritingDatasource33, InsertZorderBeforeWritingHive33, ResolveZorder}
 
@@ -29,6 +30,8 @@ class KyuubiSparkSQLCommonExtension extends (SparkSessionExtensions => Unit) {
 
 object KyuubiSparkSQLCommonExtension {
   def injectCommonExtensions(extensions: SparkSessionExtensions): Unit = {
+    extensions.injectParser((spark, parser) => new ReplaceIcebergTableParser(spark, parser))
+
     // inject zorder parser and related rules
     extensions.injectParser { case (_, parser) => new SparkKyuubiSparkSQLParser(parser) }
     extensions.injectResolutionRule(ResolveZorder)
diff --git a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/spark/sql/catalyst/parser/ReplaceIcebergTableParser.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/spark/sql/catalyst/parser/ReplaceIcebergTableParser.scala
new file mode 100644
index 00000000000..59066c81d93
--- /dev/null
+++ b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/spark/sql/catalyst/parser/ReplaceIcebergTableParser.scala
@@ -0,0 +1,155 @@
+/*
+ * 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.sql.catalyst.parser
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.{FunctionIdentifier, SQLConfHelper, TableIdentifier}
+import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation, UnresolvedTableOrView}
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.apache.spark.sql.catalyst.parser.ReplaceIcebergTableParser.ICEBERG_CATALOG
+import org.apache.spark.sql.catalyst.plans.logical.{CreateTable, CreateTableAsSelect, CreateView, InsertIntoStatement, LogicalPlan, MergeIntoTable, SubqueryAlias, UnresolvedWith}
+import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogNotFoundException, LookupCatalog, TableCatalog}
+import org.apache.spark.sql.types.{DataType, StructType}
+
+import org.apache.kyuubi.sql.KyuubiSQLConf.ICEBERG_TABLE_IMPLICIT_CONVERSION
+
+class ReplaceIcebergTableParser(sparkSession: SparkSession, delegate: ParserInterface)
+  extends ParserInterface with LookupCatalog with SQLConfHelper with Logging {
+
+  /** Creates LogicalPlan for a given SQL string. */
+  override def parsePlan(sqlText: String): LogicalPlan = {
+    val oldPlan = delegate.parsePlan(sqlText)
+    if (conf.getConf(ICEBERG_TABLE_IMPLICIT_CONVERSION)) {
+      val newPlan = resolvePlan(oldPlan)
+      newPlan
+    } else {
+      oldPlan
+    }
+  }
+
+  private def resolvePlan(plan: LogicalPlan): LogicalPlan = {
+    plan resolveOperators {
+      // insert into/overwrite
+      case i @ InsertIntoStatement(tableIdent: UnresolvedRelation, _, _, _, _, _, _) =>
+        i.copy(table = tableIdent
+          .copy(multipartIdentifier = replaceIcebergTable(tableIdent.multipartIdentifier)))
+      // CTE
+      case i @ UnresolvedWith(_, cteRelations) =>
+        i.copy(cteRelations =
+          cteRelations.map(r => (r._1, resolvePlan(r._2).asInstanceOf[SubqueryAlias])))
+      // unresolved table
+      case i @ UnresolvedRelation(multipartIdentifier, _, _) =>
+        i.copy(multipartIdentifier = replaceIcebergTable(multipartIdentifier))
+      case i @ UnresolvedTableOrView(multipartIdentifier, _, _) =>
+        i.copy(multipartIdentifier = replaceIcebergTable(multipartIdentifier))
+      // create table
+      case i @ CreateTable(UnresolvedIdentifier(nameParts, false), _, _, tableSpec, _)
+          if tableSpec.provider.isDefined && tableSpec.provider.get.toLowerCase.equals(
+            "iceberg") =>
+        val newNameParts = nameParts match {
+          case Seq(table) => Seq(ICEBERG_CATALOG, catalogManager.currentNamespace.last, table)
+          case Seq(db, table) => Seq(ICEBERG_CATALOG, db, table)
+          case _ => nameParts
+        }
+        i.copy(name = UnresolvedIdentifier(newNameParts))
+      case i @ CreateTableAsSelect(
+            UnresolvedIdentifier(nameParts, false),
+            _,
+            _,
+            tableSpec,
+            _,
+            _,
+            _)
+          if tableSpec.provider.isDefined && tableSpec.provider.get.toLowerCase.equals(
+            "iceberg") =>
+        val newNameParts = nameParts match {
+          case Seq(table) => Seq(ICEBERG_CATALOG, catalogManager.currentNamespace.last, table)
+          case Seq(db, table) => Seq(ICEBERG_CATALOG, db, table)
+          case _ => nameParts
+        }
+        i.copy(name = UnresolvedIdentifier(newNameParts))
+      case i @ CreateView(_, _, _, _, _, query, _, _) =>
+        i.copy(query = resolvePlan(query))
+      // merge into
+      case i @ MergeIntoTable(targetTable, sourceTable, _, _, _, _) =>
+        i.copy(targetTable = resolvePlan(targetTable), sourceTable = resolvePlan(sourceTable))
+    }
+  }
+
+  override protected lazy val catalogManager: CatalogManager =
+    sparkSession.sessionState.catalogManager
+
+  private def replaceIcebergTable(multipartIdentifier: Seq[String]): Seq[String] = {
+    multipartIdentifier match {
+      case SessionCatalogAndIdentifier(_, ident) =>
+        try {
+          val isIcebergTable = catalogManager.catalog(ICEBERG_CATALOG) match {
+            case tableCatalog: TableCatalog =>
+              tableCatalog.tableExists(ident)
+            case _ => false
+          }
+          if (isIcebergTable) {
+            val newIdentifier = Seq(ICEBERG_CATALOG, ident.namespace().head, ident.name())
+            logInfo(s"replace iceberg table identifier, origin: $multipartIdentifier, " +
+              s"replaced: $newIdentifier")
+            newIdentifier
+          } else {
+            multipartIdentifier
+          }
+        } catch {
+          case _: CatalogNotFoundException =>
+            multipartIdentifier
+        }
+      case _ => multipartIdentifier
+    }
+  }
+
+  override def parseExpression(sqlText: String): Expression =
+    delegate.parseExpression(sqlText)
+
+  override def parseTableIdentifier(sqlText: String): TableIdentifier =
+    delegate.parseTableIdentifier(sqlText)
+
+  override def parseFunctionIdentifier(sqlText: String): FunctionIdentifier =
+    delegate.parseFunctionIdentifier(sqlText)
+
+  override def parseMultipartIdentifier(sqlText: String): Seq[String] =
+    delegate.parseMultipartIdentifier(sqlText)
+
+  override def parseTableSchema(sqlText: String): StructType =
+    delegate.parseTableSchema(sqlText)
+
+  override def parseDataType(sqlText: String): DataType =
+    delegate.parseDataType(sqlText)
+
+  override def parseQuery(sqlText: String): LogicalPlan = {
+    val oldPlan = delegate.parseQuery(sqlText)
+    if (conf.getConf(ICEBERG_TABLE_IMPLICIT_CONVERSION)) {
+      val newPlan = resolvePlan(oldPlan)
+      newPlan
+    } else {
+      oldPlan
+    }
+  }
+}
+
+object ReplaceIcebergTableParser {
+
+  val ICEBERG_CATALOG = "iceberg_catalog"
+
+}
diff --git a/extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/ReplaceIcebergTableSuite.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/ReplaceIcebergTableSuite.scala
new file mode 100644
index 00000000000..4e8d9959376
--- /dev/null
+++ b/extensions/spark/kyuubi-extension-spark-3-5/src/test/scala/org/apache/spark/sql/ReplaceIcebergTableSuite.scala
@@ -0,0 +1,215 @@
+/*
+ * 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.sql
+
+import java.nio.file.Path
+
+import org.apache.iceberg.spark.source.SparkTable
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.catalyst.analysis.{ResolvedIdentifier, ResolvedTable}
+import org.apache.spark.sql.catalyst.plans.logical.{AppendData, CreateTable, CreateTableAsSelect, LogicalPlan, OverwriteByExpression}
+import org.apache.spark.sql.execution.command.CreateViewCommand
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.hive.execution.CreateHiveTableAsSelectCommand
+import org.apache.spark.util.Utils
+
+import org.apache.kyuubi.sql.KyuubiSQLConf
+
+class ReplaceIcebergTableSuite extends KyuubiSparkSQLExtensionTest {
+
+  protected def catalog: String = "iceberg_catalog"
+
+  protected def warehouse: Path = Utils.createTempDir().toPath
+
+  protected def extraJars: String = {
+    System.getProperty("java.class.path")
+      .split(":")
+      .filter(_.contains("iceberg-spark")).head
+  }
+
+  protected def extraConfigs = Map(
+    "spark.sql.extensions" ->
+      ("org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions," +
+        "org.apache.kyuubi.sql.KyuubiSparkSQLExtension,"),
+    s"spark.sql.catalog.$catalog" -> "org.apache.iceberg.spark.SparkCatalog",
+    s"spark.sql.catalog.$catalog.type" -> "hadoop",
+    s"spark.sql.catalog.$catalog.warehouse" -> warehouse.toString,
+    "spark.jars" -> extraJars)
+
+  override def sparkConf(): SparkConf = {
+    super.sparkConf().setAll(extraConfigs)
+  }
+
+  test("replace iceberg table") {
+    withSQLConf(KyuubiSQLConf.ICEBERG_TABLE_IMPLICIT_CONVERSION.key -> "true") {
+      spark.sql("create table default.hive_table (id int, name string)")
+
+      def checkIcebergCatalog(
+          checkSql: String,
+          pf: PartialFunction[LogicalPlan, LogicalPlan],
+          nodesWithIcebergCatalog: Int = 1): Unit = {
+        val plan = sql(checkSql).queryExecution.analyzed
+        assert(plan.collect(pf).size == nodesWithIcebergCatalog)
+      }
+
+      checkIcebergCatalog(
+        """
+          |create table iceberg_table (
+          |  id int,
+          |  name string)
+          |USING iceberg
+          |""".stripMargin,
+        {
+          case i @ CreateTable(ResolvedIdentifier(catalog, tableName), _, _, _, _)
+              if catalog.name().equals("iceberg_catalog") &&
+                tableName.name().equals("iceberg_table") => i
+        })
+
+      checkIcebergCatalog(
+        "desc iceberg_table",
+        {
+          case i @ ResolvedTable(_, _, table: SparkTable, _)
+              if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+        })
+
+      checkIcebergCatalog(
+        "select * from iceberg_table",
+        {
+          case i @ DataSourceV2Relation(table: SparkTable, _, _, _, _)
+              if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+        })
+
+      checkIcebergCatalog(
+        "insert into hive_table select * from iceberg_table",
+        {
+          case i @ DataSourceV2Relation(table: SparkTable, _, _, _, _)
+              if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+        })
+
+      checkIcebergCatalog(
+        "insert into iceberg_table values (1, 'z')",
+        {
+          case i @ AppendData(DataSourceV2Relation(table, _, _, _, _), _, _, _, _, _)
+              if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+        })
+
+      checkIcebergCatalog(
+        "insert overwrite iceberg_table values (1, 'z')",
+        {
+          case i @ OverwriteByExpression(DataSourceV2Relation(table, _, _, _, _), _, _, _, _, _, _)
+              if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+        })
+
+      checkIcebergCatalog(
+        """
+          |with t as (select * from iceberg_table)
+          |insert into hive_table select * from t
+          |""".stripMargin,
+        {
+          case i @ DataSourceV2Relation(table: SparkTable, _, _, _, _)
+              if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+        })
+
+      // TODO: https://issues.apache.org/jira/browse/SPARK-45943
+      withSQLConf("spark.sql.optimizer.runtime.rowLevelOperationGroupFilter.enabled" -> "false ") {
+        checkIcebergCatalog(
+          """
+            |MERGE INTO iceberg_table t USING (SELECT * FROM hive_table) u ON t.id = u.id
+            |WHEN MATCHED THEN UPDATE SET t.name = u.name
+            |WHEN NOT MATCHED THEN INSERT *
+            |""".stripMargin,
+          {
+            case i @ DataSourceV2Relation(table, _, _, _, _)
+                if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+          })
+      }
+
+      checkIcebergCatalog(
+        """
+          |create table iceberg_table_002
+          |USING iceberg
+          |AS select * from iceberg_table
+          |""".stripMargin,
+        {
+          case i @ CreateTableAsSelect(ResolvedIdentifier(catalog, tableName), _, query, _, _, _, _)
+              if catalog.name().equals("iceberg_catalog") &&
+                tableName.name().equals("iceberg_table_002") =>
+            val queryTable = query.collect {
+              case i @ DataSourceV2Relation(table: SparkTable, _, _, _, _)
+                  if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+            }
+            if (queryTable.size == 1) {
+              i
+            } else {
+              fail("Can not find iceberg table in query")
+            }
+        })
+
+      checkIcebergCatalog(
+        """
+          |create table hive_table_002
+          |AS select * from iceberg_table
+          |""".stripMargin,
+        {
+          case i @ CreateHiveTableAsSelectCommand(_, query, _, _) =>
+            val queryTable = query.collect {
+              case i @ DataSourceV2Relation(table: SparkTable, _, _, _, _)
+                  if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+            }
+            if (queryTable.size == 1) {
+              i
+            } else {
+              fail("Can not find iceberg table in query")
+            }
+        })
+
+      checkIcebergCatalog(
+        """
+          |UPDATE iceberg_table t
+          |SET t.name = (
+          | SELECT first(u.name) FROM hive_table u WHERE t.id = u.id
+          |)
+          |""".stripMargin,
+        {
+          case i @ DataSourceV2Relation(table, _, _, _, _)
+              if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+        })
+
+      checkIcebergCatalog(
+        """
+          |create view iceberg_view
+          |AS select * from iceberg_table
+          |""".stripMargin,
+        {
+          case i @ CreateViewCommand(_, _, _, _, _, plan, _, _, _, _, _) if plan.collect {
+                case c @ DataSourceV2Relation(table, _, _, _, _)
+                    if table.name().equals("iceberg_catalog.default.iceberg_table") => c
+              }.nonEmpty => i
+        })
+
+      withSQLConf("spark.sql.legacy.useCurrentConfigsForView" -> "true") {
+        checkIcebergCatalog(
+          "select * from iceberg_view",
+          {
+            case i @ DataSourceV2Relation(table: SparkTable, _, _, _, _)
+                if table.name().equals("iceberg_catalog.default.iceberg_table") => i
+          })
+      }
+    }
+  }
+
+}