Address comments

Author: mccheah

sql/catalyst/src/main/java/org/apache/spark/sql/catalog/v2/StagingTableCatalog.java

@@ -21,6 +21,7 @@
 
 import org.apache.spark.sql.catalog.v2.expressions.Transform;
 import org.apache.spark.sql.catalyst.analysis.NoSuchNamespaceException;
+import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
 import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException;
 import org.apache.spark.sql.sources.v2.StagedTable;
 import org.apache.spark.sql.sources.v2.SupportsWrite;
@@ -87,10 +88,11 @@ StagedTable stageCreate(
    * can decide whether to move forward with the table replacement anyways or abort the commit
    * operation.
    * <p>
-   * If the table does not exist, committing the staged changes should fail. This differs from the
-   * semantics of {@link #stageCreateOrReplace(Identifier, StructType, Transform[], Map)}, which
-   * should create the table in the data source if the table does not exist at the time of
-   * committing the operation.
+   * If the table does not exist, committing the staged changes should fail with
+   * {@link NoSuchTableException}. This differs from the semantics of
+   * {@link #stageCreateOrReplace(Identifier, StructType, Transform[], Map)}, which should create
+   * the table in the data source if the table does not exist at the time of committing the
+   * operation.
    *
    * @param ident a table identifier
    * @param schema the schema of the new table, as a struct type
@@ -99,12 +101,13 @@ StagedTable stageCreate(
    * @return metadata for the new table
    * @throws UnsupportedOperationException If a requested partition transform is not supported
    * @throws NoSuchNamespaceException If the identifier namespace does not exist (optional)
+   * @throws NoSuchTableException If the table does not exist
    */
   StagedTable stageReplace(
       Identifier ident,
       StructType schema,
       Transform[] partitions,
-      Map<String, String> properties) throws NoSuchNamespaceException;
+      Map<String, String> properties) throws NoSuchNamespaceException, NoSuchTableException;
 
   /**
    * Stage the creation or replacement of a table, preparing it to be committed into the metastore

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

@@ -22,7 +22,8 @@ import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalog.v2.Identifier
 
 class CannotReplaceMissingTableException(
-    tableIdentifier: Identifier)
+    tableIdentifier: Identifier,
+    cause: Option[Throwable] = None)
   extends AnalysisException(
     s"Table $tableIdentifier cannot be replaced as it did not exist." +
-    s" Use CREATE OR REPLACE TABLE to create the table.")
+    s" Use CREATE OR REPLACE TABLE to create the table.", cause = cause)

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

@@ -2308,7 +2308,7 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
    *
    * Expected format:
    * {{{
-   *   REPLACE TABLE [IF NOT EXISTS] [db_name.]table_name
+   *   [CREATE OR] REPLACE TABLE [db_name.]table_name
    *   USING table_provider
    *   replace_table_clauses
    *   [[AS] select_statement];

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ReplaceTableExec.scala

@@ -23,7 +23,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalog.v2.{Identifier, StagingTableCatalog, TableCatalog}
 import org.apache.spark.sql.catalog.v2.expressions.Transform
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.analysis.CannotReplaceMissingTableException
+import org.apache.spark.sql.catalyst.analysis.{CannotReplaceMissingTableException, NoSuchTableException}
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.execution.LeafExecNode
 import org.apache.spark.sql.sources.v2.StagedTable
@@ -32,19 +32,19 @@ import org.apache.spark.util.Utils
 
 case class ReplaceTableExec(
     catalog: TableCatalog,
-    identifier: Identifier,
+    ident: Identifier,
     tableSchema: StructType,
     partitioning: Seq[Transform],
     tableProperties: Map[String, String],
     orCreate: Boolean) extends LeafExecNode {
 
   override protected def doExecute(): RDD[InternalRow] = {
-    if (catalog.tableExists(identifier)) {
-      catalog.dropTable(identifier)
+    if (catalog.tableExists(ident)) {
+      catalog.dropTable(ident)
     } else if (!orCreate) {
-      throw new CannotReplaceMissingTableException(identifier)
+      throw new CannotReplaceMissingTableException(ident)
     }
-    catalog.createTable(identifier, tableSchema, partitioning.toArray, tableProperties.asJava)
+    catalog.createTable(ident, tableSchema, partitioning.toArray, tableProperties.asJava)
     sqlContext.sparkContext.parallelize(Seq.empty, 1)
   }
 
@@ -60,18 +60,17 @@ case class AtomicReplaceTableExec(
     orCreate: Boolean) extends LeafExecNode {
 
   override protected def doExecute(): RDD[InternalRow] = {
-    val staged = if (catalog.tableExists(identifier)) {
-      catalog.stageReplace(
-        identifier,
-        tableSchema,
-        partitioning.toArray,
-        tableProperties.asJava)
-    } else if (orCreate) {
-      catalog.stageCreate(
-        identifier,
-        tableSchema,
-        partitioning.toArray,
-        tableProperties.asJava)
+    val staged = if (orCreate) {
+      catalog.stageCreateOrReplace(
+        identifier, tableSchema, partitioning.toArray, tableProperties.asJava)
+    } else if (catalog.tableExists(identifier)) {
+      try {
+        catalog.stageReplace(
+          identifier, tableSchema, partitioning.toArray, tableProperties.asJava)
+      } catch {
+        case e: NoSuchTableException =>
+          throw new CannotReplaceMissingTableException(identifier, Some(e))
+      }
     } else {
       throw new CannotReplaceMissingTableException(identifier)
     }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala

@@ -29,7 +29,7 @@ import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalog.v2.{Identifier, StagingTableCatalog, TableCatalog}
 import org.apache.spark.sql.catalog.v2.expressions.Transform
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.analysis.{CannotReplaceMissingTableException, TableAlreadyExistsException}
+import org.apache.spark.sql.catalyst.analysis.{CannotReplaceMissingTableException, NoSuchTableException, TableAlreadyExistsException}
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
@@ -213,21 +213,21 @@ case class AtomicReplaceTableAsSelectExec(
     orCreate: Boolean) extends AtomicTableWriteExec {
 
   override protected def doExecute(): RDD[InternalRow] = {
-    val stagedTable = if (catalog.tableExists(ident)) {
-      if (orCreate) {
-        catalog.stageCreateOrReplace(
-          ident, query.schema, partitioning.toArray, properties.asJava)
-      } else {
+    val staged = if (orCreate) {
+      catalog.stageCreateOrReplace(
+        ident, query.schema, partitioning.toArray, properties.asJava)
+    } else if (catalog.tableExists(ident)) {
+      try {
         catalog.stageReplace(
           ident, query.schema, partitioning.toArray, properties.asJava)
+      } catch {
+        case e: NoSuchTableException =>
+          throw new CannotReplaceMissingTableException(ident, Some(e))
       }
-    } else if (orCreate) {
-      catalog.stageCreateOrReplace(
-        ident, query.schema, partitioning.toArray, properties.asJava)
     } else {
       throw new CannotReplaceMissingTableException(ident)
     }
-    writeToStagedTable(stagedTable, writeOptions, ident)
+    writeToStagedTable(staged, writeOptions, ident)
   }
 }