[SPARK-27322][SQL] DataSourceV2: Select from multiple catalogs

Author: jzhuge

sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4

@@ -397,7 +397,7 @@ queryTerm
 
 queryPrimary
     : querySpecification                                                    #queryPrimaryDefault
-    | TABLE tableIdentifier                                                 #table
+    | TABLE multipartIdentifier                                             #table
     | inlineTable                                                           #inlineTableDefault1
     | '(' queryNoWith  ')'                                                  #subquery
     ;
@@ -536,7 +536,7 @@ identifierComment
     ;
 
 relationPrimary
-    : tableIdentifier sample? tableAlias      #tableName
+    : multipartIdentifier sample? tableAlias  #tableName
     | '(' queryNoWith ')' sample? tableAlias  #aliasedQuery
     | '(' relation ')' sample? tableAlias     #aliasedRelation
     | inlineTable                             #inlineTableDefault2

sql/catalyst/src/main/scala/org/apache/spark/sql/catalog/v2/TableIdentifierHelper.scala

@@ -0,0 +1,37 @@
+/*
+ * 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.catalog.v2
+
+import org.apache.spark.sql.catalyst.{CatalogTableIdentifier, TableIdentifier, TableIdentifierLike}
+
+/**
+ * Resolve multipart table identifier to [[CatalogTableIdentifier]] or [[TableIdentifier]].
+ */
+trait TableIdentifierHelper extends LookupCatalog {
+  import CatalogV2Implicits._
+
+  implicit class TableIdentifierHelper(parts: Seq[String]) {
+    def asCatalogTableIdentifier: TableIdentifierLike = parts match {
+      case CatalogObjectIdentifier(Some(catalog), ident) =>
+        CatalogTableIdentifier(catalog.asTableCatalog, ident)
+
+      case AsTableIdentifier(tableIdentifier) =>
+        tableIdentifier
+    }
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -224,9 +224,8 @@ class Analyzer(
 
     def substituteCTE(plan: LogicalPlan, cteRelations: Seq[(String, LogicalPlan)]): LogicalPlan = {
       plan resolveOperatorsDown {
-        case u: UnresolvedRelation =>
-          cteRelations.find(x => resolver(x._1, u.tableIdentifier.table))
-            .map(_._2).getOrElse(u)
+        case u @ UnresolvedRelation(TableIdentifier(table, _)) =>
+          cteRelations.find(x => resolver(x._1, table)).map(_._2).getOrElse(u)
         case other =>
           // This cannot be done in ResolveSubquery because ResolveSubquery does not know the CTE.
           other transformExpressions {
@@ -718,7 +717,7 @@ class Analyzer(
             u.failAnalysis(s"Inserting into a view is not allowed. View: ${v.desc.identifier}.")
           case other => i.copy(table = other)
         }
-      case u: UnresolvedRelation => resolveRelation(u)
+      case u @ UnresolvedRelation(_: TableIdentifier) => resolveRelation(u)
     }
 
     // Look up the table with the given name from catalog. The database we used is decided by the

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveHints.scala

@@ -22,6 +22,7 @@ import java.util.Locale
 import scala.collection.mutable
 
 import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.expressions.IntegerLiteral
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
@@ -71,17 +72,18 @@ object ResolveHints {
 
       val newNode = CurrentOrigin.withOrigin(plan.origin) {
         plan match {
-          case ResolvedHint(u: UnresolvedRelation, hint)
-              if relations.exists(resolver(_, u.tableIdentifier.table)) =>
-            relations.remove(u.tableIdentifier.table)
+          case ResolvedHint(u @ UnresolvedRelation(TableIdentifier(table, _)), hint)
+              if relations.exists(resolver(_, table)) =>
+            relations.remove(table)
             ResolvedHint(u, createHintInfo(hintName).merge(hint, handleOverriddenHintInfo))
           case ResolvedHint(r: SubqueryAlias, hint)
               if relations.exists(resolver(_, r.alias)) =>
             relations.remove(r.alias)
             ResolvedHint(r, createHintInfo(hintName).merge(hint, handleOverriddenHintInfo))
 
-          case u: UnresolvedRelation if relations.exists(resolver(_, u.tableIdentifier.table)) =>
-            relations.remove(u.tableIdentifier.table)
+          case UnresolvedRelation(TableIdentifier(table, _))
+            if relations.exists(resolver(_, table)) =>
+            relations.remove(table)
             ResolvedHint(plan, createHintInfo(hintName))
           case r: SubqueryAlias if relations.exists(resolver(_, r.alias)) =>
             relations.remove(r.alias)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/unresolved.scala

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.catalyst.analysis
 
 import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.catalyst.{FunctionIdentifier, InternalRow, TableIdentifier}
+import org.apache.spark.sql.catalyst.{FunctionIdentifier, InternalRow, TableIdentifier, TableIdentifierLike}
 import org.apache.spark.sql.catalyst.errors.TreeNodeException
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode}
@@ -38,11 +38,13 @@ class UnresolvedException[TreeType <: TreeNode[_]](tree: TreeType, function: Str
 /**
  * Holds the name of a relation that has yet to be looked up in a catalog.
  *
- * @param tableIdentifier table name
+ * @param table table name
  */
-case class UnresolvedRelation(tableIdentifier: TableIdentifier)
+case class UnresolvedRelation(table: TableIdentifierLike)
   extends LeafNode {
 
+  def tableIdentifier: TableIdentifier = table.tableIdentifier
+
   /** Returns a `.` separated name for this relation. */
   def tableName: String = tableIdentifier.unquotedString
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/identifiers.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.catalyst
 
+import org.apache.spark.sql.catalog.v2.{Identifier, TableCatalog}
+
 /**
  * An identifier that optionally specifies a database.
  *
@@ -64,14 +66,37 @@ object AliasIdentifier {
   def apply(identifier: String): AliasIdentifier = new AliasIdentifier(identifier)
 }
 
+/**
+ * An interface to ease transition from [[TableIdentifier]] to [[CatalogTableIdentifier]].
+ */
+sealed trait TableIdentifierLike {
+  def tableIdentifier: TableIdentifier
+}
+
+/**
+ * A data source V2 table identifier.
+ *
+ * @param catalog a catalog plugin
+ * @param ident an object identifier
+ */
+case class CatalogTableIdentifier(catalog: TableCatalog, ident: Identifier)
+  extends TableIdentifierLike {
+
+  override def tableIdentifier: TableIdentifier =
+    throw new UnsupportedOperationException(
+      s"$this should not be used on non-DSv2 code path")
+}
+
 /**
  * Identifies a table in a database.
  * If `database` is not defined, the current database is used.
  * When we register a permanent function in the FunctionRegistry, we use
  * unquotedString as the function name.
  */
 case class TableIdentifier(table: String, database: Option[String])
-  extends IdentifierWithDatabase {
+  extends IdentifierWithDatabase with TableIdentifierLike {
+
+  override def tableIdentifier: TableIdentifier = this
 
   override val identifier: String = table
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -29,6 +29,7 @@ import org.antlr.v4.runtime.tree.{ParseTree, RuleNode, TerminalNode}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalog.v2
+import org.apache.spark.sql.catalog.v2.{CatalogNotFoundException, CatalogPlugin, TableIdentifierHelper}
 import org.apache.spark.sql.catalog.v2.expressions.{ApplyTransform, BucketTransform, DaysTransform, FieldReference, HoursTransform, IdentityTransform, LiteralValue, MonthsTransform, Transform, YearsTransform}
 import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier}
 import org.apache.spark.sql.catalyst.analysis._
@@ -49,11 +50,15 @@ import org.apache.spark.util.random.RandomSampler
  * The AstBuilder converts an ANTLR4 ParseTree into a catalyst Expression, LogicalPlan or
  * TableIdentifier.
  */
-class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging {
+class AstBuilder(conf: SQLConf)
+  extends SqlBaseBaseVisitor[AnyRef] with Logging with TableIdentifierHelper {
   import ParserUtils._
 
   def this() = this(new SQLConf())
 
+  override protected def lookupCatalog(name: String): CatalogPlugin =
+    throw new CatalogNotFoundException("No catalog lookup function")
+
   protected def typedVisit[T](ctx: ParseTree): T = {
     ctx.accept(this).asInstanceOf[T]
   }
@@ -844,14 +849,15 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
    * }}}
    */
   override def visitTable(ctx: TableContext): LogicalPlan = withOrigin(ctx) {
-    UnresolvedRelation(visitTableIdentifier(ctx.tableIdentifier))
+    val tableId = visitMultipartIdentifier(ctx.multipartIdentifier).asCatalogTableIdentifier
+    UnresolvedRelation(tableId)
   }
 
   /**
    * Create an aliased table reference. This is typically used in FROM clauses.
    */
   override def visitTableName(ctx: TableNameContext): LogicalPlan = withOrigin(ctx) {
-    val tableId = visitTableIdentifier(ctx.tableIdentifier)
+    val tableId = visitMultipartIdentifier(ctx.multipartIdentifier).asCatalogTableIdentifier
     val table = mayApplyAliasPlan(ctx.tableAlias, UnresolvedRelation(tableId))
     table.optionalMap(ctx.sample)(withSample)
   }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/TableIdentifierParserSuite.scala

@@ -17,13 +17,28 @@
 package org.apache.spark.sql.catalyst.parser
 
 import org.apache.spark.SparkFunSuite
-import org.apache.spark.sql.catalyst.TableIdentifier
+import org.apache.spark.sql.catalog.v2.{CatalogNotFoundException, CatalogPlugin, Identifier, TableIdentifierHelper, TestTableCatalog}
+import org.apache.spark.sql.catalyst.{CatalogTableIdentifier, TableIdentifier}
 import org.apache.spark.sql.catalyst.plans.SQLHelper
 import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
 
-class TableIdentifierParserSuite extends SparkFunSuite with SQLHelper {
+class TableIdentifierParserSuite extends SparkFunSuite with SQLHelper with TableIdentifierHelper {
   import CatalystSqlParser._
 
+  private val testCat = new TestTableCatalog {
+    initialize("testcat", CaseInsensitiveStringMap.empty())
+  }
+
+  private def findCatalog(name: String): CatalogPlugin = name match {
+    case "testcat" =>
+      testCat
+    case _ =>
+      throw new CatalogNotFoundException(s"$name not found")
+  }
+
+  override protected def lookupCatalog(name: String): CatalogPlugin = findCatalog(name)
+
   // Add "$elem$", "$value$" & "$key$"
   // It is recommended to list them in alphabetical order.
   val hiveNonReservedKeyword = Array(
@@ -699,4 +714,11 @@ class TableIdentifierParserSuite extends SparkFunSuite with SQLHelper {
     val complexName2 = TableIdentifier("x``y", Some("d``b"))
     assert(complexName2 === parseTableIdentifier(complexName2.quotedString))
   }
+
+  test("multipart table identifier") {
+    assert(parseMultipartIdentifier("testcat.v2tbl").asCatalogTableIdentifier ===
+      CatalogTableIdentifier(testCat, Identifier.of(Array.empty, "v2tbl")))
+    assert(parseMultipartIdentifier("db.tbl").asCatalogTableIdentifier ===
+      TableIdentifier("tbl", Some("db")))
+  }
 }

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

@@ -33,7 +33,7 @@ import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
 import org.apache.spark.scheduler.{SparkListener, SparkListenerApplicationEnd}
 import org.apache.spark.sql.catalog.Catalog
-import org.apache.spark.sql.catalog.v2.{CatalogPlugin, Catalogs}
+import org.apache.spark.sql.catalog.v2.{CatalogPlugin, Catalogs, TableIdentifierHelper}
 import org.apache.spark.sql.catalyst._
 import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
 import org.apache.spark.sql.catalyst.encoders._
@@ -82,7 +82,7 @@ class SparkSession private(
     @transient private val existingSharedState: Option[SharedState],
     @transient private val parentSessionState: Option[SessionState],
     @transient private[sql] val extensions: SparkSessionExtensions)
-  extends Serializable with Closeable with Logging { self =>
+  extends Serializable with Closeable with Logging with TableIdentifierHelper { self =>
 
   // The call site where this SparkSession was constructed.
   private val creationSite: CallSite = Utils.getCallSite()
@@ -613,6 +613,8 @@ class SparkSession private(
     catalogs.getOrElseUpdate(name, Catalogs.load(name, sessionState.conf))
   }
 
+  override def lookupCatalog(name: String): CatalogPlugin = catalog(name)
+
   /**
    * Returns the specified table/view as a `DataFrame`.
    *
@@ -624,10 +626,10 @@ class SparkSession private(
    * @since 2.0.0
    */
   def table(tableName: String): DataFrame = {
-    table(sessionState.sqlParser.parseTableIdentifier(tableName))
+    table(sessionState.sqlParser.parseMultipartIdentifier(tableName).asCatalogTableIdentifier)
   }
 
-  private[sql] def table(tableIdent: TableIdentifier): DataFrame = {
+  private[sql] def table(tableIdent: TableIdentifierLike): DataFrame = {
     Dataset.ofRows(self, UnresolvedRelation(tableIdent))
   }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/command/views.scala

@@ -190,10 +190,10 @@ case class CreateViewCommand(
       // package (e.g., HiveGenericUDF).
       child.collect {
         // Disallow creating permanent views based on temporary views.
-        case s: UnresolvedRelation
-          if sparkSession.sessionState.catalog.isTemporaryTable(s.tableIdentifier) =>
+        case UnresolvedRelation(tableIdentifier: TableIdentifier)
+          if sparkSession.sessionState.catalog.isTemporaryTable(tableIdentifier) =>
           throw new AnalysisException(s"Not allowed to create a permanent view $name by " +
-            s"referencing a temporary view ${s.tableIdentifier}")
+            s"referencing a temporary view ${tableIdentifier}")
         case other if !other.resolved => other.expressions.flatMap(_.collect {
           // Disallow creating permanent views based on temporary UDFs.
           case e: UnresolvedFunction