[SPARK-32512][SQL] add alter table add/drop partition command for datasourcev2

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

@@ -17,10 +17,12 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
+import scala.collection.JavaConverters._
+
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogPlugin, LookupCatalog, SupportsNamespaces, TableCatalog, TableChange}
+import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogPlugin, LookupCatalog, SupportsNamespaces, SupportsPartitionManagement, TableCatalog, TableChange}
 
 /**
  * Resolves catalogs from the multi-part identifiers in SQL statements, and convert the statements
@@ -30,6 +32,7 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
   extends Rule[LogicalPlan] with LookupCatalog {
   import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._
   import org.apache.spark.sql.connector.catalog.CatalogV2Util._
+  import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._
 
   override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
     case AlterTableAddColumnsStatement(
@@ -228,6 +231,30 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
 
     case ShowCurrentNamespaceStatement() =>
       ShowCurrentNamespace(catalogManager)
+
+    case c @ AlterTableAddPartitionStatement(
+        ResolvedTable(_, _, table: SupportsPartitionManagement), _, _) =>
+      val partitions = c.partitionSpecsAndLocs.map { case (spec, location) =>
+        val tableProperties = table.properties().asScala.toMap
+        val partParams =
+          location.map(locationUri => tableProperties + ("location" -> locationUri))
+            .getOrElse(tableProperties)
+        (spec.asPartitionIdentifier(table.partitionSchema()), partParams)
+      }
+      AlterTableAddPartition(
+        table,
+        partitions,
+        c.ifNotExists)
+
+    case c @ AlterTableDropPartitionStatement(
+        ResolvedTable(_, _, table: SupportsPartitionManagement), _, _, _, _) =>
+      if (c.purge) {
+        logWarning("PURGE won't take effect here, please put it in table properties")
+      }
+      AlterTableDropPartition(
+        table,
+        c.specs.map(_.asPartitionIdentifier(table.partitionSchema())),
+        c.ifExists)
   }
 
   object NonSessionCatalogAndTable {

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

@@ -3418,7 +3418,7 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
       spec -> location
     }
     AlterTableAddPartitionStatement(
-      visitMultipartIdentifier(ctx.multipartIdentifier),
+      UnresolvedTableOrView(visitMultipartIdentifier(ctx.multipartIdentifier)),
       specsAndLocs.toSeq,
       ctx.EXISTS != null)
   }
@@ -3458,7 +3458,7 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
       operationNotAllowed("ALTER VIEW ... DROP PARTITION", ctx)
     }
     AlterTableDropPartitionStatement(
-      visitMultipartIdentifier(ctx.multipartIdentifier),
+      UnresolvedTableOrView(visitMultipartIdentifier(ctx.multipartIdentifier)),
       ctx.partitionSpec.asScala.map(visitNonOptionalPartitionSpec).toSeq,
       ifExists = ctx.EXISTS != null,
       purge = ctx.PURGE != null,

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/statements.scala

@@ -221,7 +221,7 @@ case class AlterTableRecoverPartitionsStatement(
  * ALTER TABLE ... ADD PARTITION command, as parsed from SQL
  */
 case class AlterTableAddPartitionStatement(
-    tableName: Seq[String],
+    child: LogicalPlan,
     partitionSpecsAndLocs: Seq[(TablePartitionSpec, Option[String])],
     ifNotExists: Boolean) extends ParsedStatement
 
@@ -237,7 +237,7 @@ case class AlterTableRenamePartitionStatement(
  * ALTER TABLE ... DROP PARTITION command, as parsed from SQL
  */
 case class AlterTableDropPartitionStatement(
-    tableName: Seq[String],
+    child: LogicalPlan,
     specs: Seq[TablePartitionSpec],
     ifExists: Boolean,
     purge: Boolean,

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.catalyst.plans.logical
 
+import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.{NamedRelation, UnresolvedException}
 import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec
 import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Expression, Unevaluable}
@@ -551,3 +552,31 @@ case class ShowFunctions(
     pattern: Option[String]) extends Command {
   override def children: Seq[LogicalPlan] = child.toSeq
 }
+
+/**
+ * The logical plan of the ALTER TABLE ADD PARTITION command that works for v2 tables.
+ *
+ * The syntax of this command is:
+ * {{{
+ *     ALTER TABLE table ADD [IF NOT EXISTS]
+ *                 PARTITION spec1 [LOCATION 'loc1'][, PARTITION spec2 [LOCATION 'loc2'], ...];
+ * }}}
+ */
+case class AlterTableAddPartition(
+    table: SupportsPartitionManagement,
+    parts: Seq[(InternalRow, Map[String, String])],
+    ignoreIfExists: Boolean) extends Command
+
+/**
+ * The logical plan of the ALTER TABLE DROP PARTITION command that works for v2 tables.
+ * This may remove the data and metadata for this partition.
+ *
+ * The syntax of this command is:
+ * {{{
+ *     ALTER TABLE table DROP [IF EXISTS] PARTITION spec1[, PARTITION spec2, ...];
+ * }}}
+ */
+case class AlterTableDropPartition(
+    table: SupportsPartitionManagement,
+    partIdents: Seq[InternalRow],
+    ignoreIfNotExists: Boolean) extends Command

sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Implicits.scala

@@ -20,7 +20,10 @@ package org.apache.spark.sql.execution.datasources.v2
 import scala.collection.JavaConverters._
 
 import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.connector.catalog.{SupportsDelete, SupportsRead, SupportsWrite, Table, TableCapability}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec
+import org.apache.spark.sql.connector.catalog.{SupportsAtomicPartitionManagement, SupportsDelete, SupportsPartitionManagement, SupportsRead, SupportsWrite, Table, TableCapability}
+import org.apache.spark.sql.types.{ByteType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, StructType}
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 
 object DataSourceV2Implicits {
@@ -52,6 +55,26 @@ object DataSourceV2Implicits {
       }
     }
 
+    def asPartitionable: SupportsPartitionManagement = {
+      table match {
+        case support: SupportsPartitionManagement =>
+          support
+        case _ =>
+          throw new AnalysisException(
+            s"Table does not support partition management: ${table.name}")
+      }
+    }
+
+    def asAtomicPartitionable: SupportsAtomicPartitionManagement = {
+      table match {
+        case support: SupportsAtomicPartitionManagement =>
+          support
+        case _ =>
+          throw new AnalysisException(
+            s"Table does not support atomic partition management: ${table.name}")
+      }
+    }
+
     def supports(capability: TableCapability): Boolean = table.capabilities.contains(capability)
 
     def supportsAny(capabilities: TableCapability*): Boolean = capabilities.exists(supports)
@@ -62,4 +85,42 @@ object DataSourceV2Implicits {
       new CaseInsensitiveStringMap(options.asJava)
     }
   }
+
+  implicit class TablePartitionSpecHelper(partSpec: TablePartitionSpec) {
+    def asPartitionIdentifier(partSchema: StructType): InternalRow = {
+      val conflictKeys = partSpec.keys.toSeq.diff(partSchema.map(_.name))
+      if (conflictKeys.nonEmpty) {
+        throw new AnalysisException(s"Partition key ${conflictKeys.mkString(",")} not exists")
+      }
+
+      val partValues = partSchema.map { part =>
+        val partValue = partSpec.get(part.name).orNull
+        if (partValue == null) {
+          null
+        } else {
+          // TODO: Support other datatypes, such as DateType
+          part.dataType match {
+            case _: ByteType =>
+              partValue.toByte
+            case _: ShortType =>
+              partValue.toShort
+            case _: IntegerType =>
+              partValue.toInt
+            case _: LongType =>
+              partValue.toLong
+            case _: FloatType =>
+              partValue.toFloat
+            case _: DoubleType =>
+              partValue.toDouble
+            case _: StringType =>
+              partValue
+            case _ =>
+              throw new AnalysisException(
+                s"Type ${part.dataType.typeName} is not supported for partition.")
+          }
+        }
+      }
+      InternalRow.fromSeq(partValues)
+    }
+  }
 }

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

@@ -1708,13 +1708,13 @@ class DDLParserSuite extends AnalysisTest {
     val parsed2 = parsePlan(sql2)
 
     val expected1 = AlterTableAddPartitionStatement(
-      Seq("a", "b", "c"),
+      UnresolvedTableOrView(Seq("a", "b", "c")),
       Seq(
         (Map("dt" -> "2008-08-08", "country" -> "us"), Some("location1")),
         (Map("dt" -> "2009-09-09", "country" -> "uk"), None)),
       ifNotExists = true)
     val expected2 = AlterTableAddPartitionStatement(
-      Seq("a", "b", "c"),
+      UnresolvedTableOrView(Seq("a", "b", "c")),
       Seq((Map("dt" -> "2008-08-08"), Some("loc"))),
       ifNotExists = false)
 
@@ -1781,7 +1781,7 @@ class DDLParserSuite extends AnalysisTest {
     assertUnsupported(sql2_view)
 
     val expected1_table = AlterTableDropPartitionStatement(
-      Seq("table_name"),
+      UnresolvedTableOrView(Seq("table_name")),
       Seq(
         Map("dt" -> "2008-08-08", "country" -> "us"),
         Map("dt" -> "2009-09-09", "country" -> "uk")),
@@ -1797,7 +1797,7 @@ class DDLParserSuite extends AnalysisTest {
 
     val sql3_table = "ALTER TABLE a.b.c DROP IF EXISTS PARTITION (ds='2017-06-10')"
     val expected3_table = AlterTableDropPartitionStatement(
-      Seq("a", "b", "c"),
+      UnresolvedTableOrView(Seq("a", "b", "c")),
       Seq(Map("ds" -> "2017-06-10")),
       ifExists = true,
       purge = false,

sql/catalyst/src/test/scala/org/apache/spark/sql/connector/InMemoryPartitionTable.scala

@@ -92,4 +92,8 @@ class InMemoryPartitionTable(
 
   override def partitionExists(ident: InternalRow): Boolean =
     memoryTablePartitions.containsKey(ident)
+
+  def clearPartitions(): Unit = {
+    memoryTablePartitions.clear()
+  }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/connector/InMemoryPartitionTableCatalog.scala

@@ -0,0 +1,47 @@
+/*
+ * 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.connector
+
+import java.util
+
+import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException
+import org.apache.spark.sql.connector.catalog.{CatalogV2Implicits, Identifier, Table}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.StructType
+
+class InMemoryPartitionTableCatalog extends InMemoryTableCatalog {
+  import CatalogV2Implicits._
+
+  override def createTable(
+      ident: Identifier,
+      schema: StructType,
+      partitions: Array[Transform],
+      properties: util.Map[String, String]): Table = {
+    if (tables.containsKey(ident)) {
+      throw new TableAlreadyExistsException(ident)
+    }
+
+    InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties)
+
+    val table = new InMemoryAtomicPartitionTable(
+      s"$name.${ident.quoted}", schema, partitions, properties)
+    tables.put(ident, table)
+    namespaces.putIfAbsent(ident.namespace.toList, Map())
+    table
+  }
+}

sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala

@@ -22,7 +22,7 @@ import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier}
 import org.apache.spark.sql.catalyst.catalog.{BucketSpec, CatalogStorageFormat, CatalogTable, CatalogTableType, CatalogUtils}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
-import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogPlugin, CatalogV2Util, LookupCatalog, SupportsNamespaces, TableCatalog, TableChange, V1Table}
+import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogPlugin, CatalogV2Util, LookupCatalog, SupportsNamespaces, SupportsPartitionManagement, TableCatalog, TableChange, V1Table}
 import org.apache.spark.sql.connector.expressions.Transform
 import org.apache.spark.sql.execution.command._
 import org.apache.spark.sql.execution.datasources.{CreateTable, DataSource, RefreshTable}
@@ -497,10 +497,11 @@ class ResolveSessionCatalog(
         v1TableName.asTableIdentifier,
         "ALTER TABLE RECOVER PARTITIONS")
 
-    case AlterTableAddPartitionStatement(tbl, partitionSpecsAndLocs, ifNotExists) =>
-      val v1TableName = parseV1Table(tbl, "ALTER TABLE ADD PARTITION")
+    case AlterTableAddPartitionStatement(
+        r @ ResolvedTable(_, _, _: V1Table), partitionSpecsAndLocs, ifNotExists)
+        if isSessionCatalog(r.catalog) =>
       AlterTableAddPartitionCommand(
-        v1TableName.asTableIdentifier,
+        r.identifier.asTableIdentifier,
         partitionSpecsAndLocs,
         ifNotExists)
 
@@ -511,10 +512,11 @@ class ResolveSessionCatalog(
         from,
         to)
 
-    case AlterTableDropPartitionStatement(tbl, specs, ifExists, purge, retainData) =>
-      val v1TableName = parseV1Table(tbl, "ALTER TABLE DROP PARTITION")
+    case AlterTableDropPartitionStatement(
+        r @ ResolvedTable(_, _, _: V1Table), specs, ifExists, purge, retainData)
+        if isSessionCatalog(r.catalog) =>
       AlterTableDropPartitionCommand(
-        v1TableName.asTableIdentifier,
+        r.identifier.asTableIdentifier,
         specs,
         ifExists,
         purge,