11

Author: Zouxxyy

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

@@ -1364,6 +1364,7 @@ case class DropPartitions(
     parts: Seq[PartitionSpec],
     ifExists: Boolean,
     purge: Boolean) extends V2PartitionCommand {
+  override def allowPartialPartitionSpec: Boolean = true
   override protected def withNewChildInternal(newChild: LogicalPlan): DropPartitions =
     copy(table = newChild)
 }

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

@@ -17,6 +17,8 @@
 
 package org.apache.spark.sql.execution.datasources.v2
 
+import scala.collection.mutable
+
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.{NoSuchPartitionsException, ResolvedPartitionSpec}
 import org.apache.spark.sql.catalyst.expressions.Attribute
@@ -37,15 +39,23 @@ case class DropPartitionExec(
   override def output: Seq[Attribute] = Seq.empty
 
   override protected def run(): Seq[InternalRow] = {
-    val (existsPartIdents, notExistsPartIdents) =
-      partSpecs.map(_.ident).partition(table.partitionExists)
+    val existsPartIdents = mutable.Set[InternalRow]()
+    val notExistsPartIdents = mutable.Set[InternalRow]()
+    partSpecs.foreach { partSpec =>
+      val partIdents = table.listPartitionIdentifiers(partSpec.names.toArray, partSpec.ident)
+      if (partIdents.nonEmpty) {
+        existsPartIdents.addAll(partIdents)
+      } else {
+        notExistsPartIdents.add(partSpec.ident)
+      }
+    }
 
     if (notExistsPartIdents.nonEmpty && !ignoreIfNotExists) {
       throw new NoSuchPartitionsException(
-        table.name(), notExistsPartIdents, table.partitionSchema())
+        table.name(), notExistsPartIdents.toSeq, table.partitionSchema())
     }
 
-    val isTableAltered = existsPartIdents match {
+    val isTableAltered = existsPartIdents.toSeq match {
       case Seq() => false // Nothing will be done
       case Seq(partIdent) =>
         if (purge) table.purgePartition(partIdent) else table.dropPartition(partIdent)

sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala

@@ -3042,6 +3042,37 @@ class DataSourceV2SQLSuiteV1Filter
     }
   }
 
+  test("Drop multi-level partition table with partial partition spec") {
+    val t = "testpart.tbl"
+    withTable(t) {
+      sql(s"CREATE TABLE $t (id string) USING foo PARTITIONED BY (p1 int, p2 int, p3 int)")
+      sql(s"INSERT INTO $t VALUES ('a', 1, 11, 111), ('b', 1, 22, 222), ('c', 2, 22, 333)")
+
+      checkAnswer(
+        sql(s"SHOW PARTITIONS $t"),
+        Seq(Row("p1=1/p2=11/p3=111"), Row("p1=1/p2=22/p3=222"), Row("p1=2/p2=22/p3=333")))
+
+      sql(s"ALTER TABLE $t DROP PARTITION (p1=1)")
+      checkAnswer(
+        sql(s"SHOW PARTITIONS $t"),
+        Seq(Row("p1=2/p2=22/p3=333")))
+
+      checkError(exception = intercept[AnalysisException] {
+        sql(s"ALTER TABLE $t DROP PARTITION (p1=4)")
+      },
+        condition = "PARTITIONS_NOT_FOUND",
+        sqlState = Some("428FT"),
+        parameters = Map(
+          "partitionList" -> "PARTITION (`p1` = 4)",
+          "tableName" -> t.split("\\.").map(part => s"`$part`").mkString(".")))
+
+      sql(s"ALTER TABLE $t DROP IF EXISTS PARTITION (p1=4)")
+      checkAnswer(
+        sql(s"SHOW PARTITIONS $t"),
+        Seq(Row("p1=2/p2=22/p3=333")))
+    }
+  }
+
   test("Check HasPartitionStatistics from InMemoryPartitionTable") {
     val t = "testpart.tbl"
     withTable(t) {