[SPARK-14922][SPARK-17732][SPARK-23866][SQL] Support partition filters in ALTER TABLE DROP PARTITION

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

@@ -114,10 +114,10 @@ statement
         partitionSpec+                                                 #addTablePartition
     | ALTER TABLE tableIdentifier
         from=partitionSpec RENAME TO to=partitionSpec                  #renameTablePartition
-    | ALTER TABLE tableIdentifier
-        DROP (IF EXISTS)? partitionSpec (',' partitionSpec)* PURGE?    #dropTablePartitions
-    | ALTER VIEW tableIdentifier
-        DROP (IF EXISTS)? partitionSpec (',' partitionSpec)*           #dropTablePartitions
+    | ALTER TABLE tableIdentifier DROP (IF EXISTS)?
+        dropPartitionSpec (',' dropPartitionSpec)* PURGE?              #dropTablePartitions
+    | ALTER VIEW tableIdentifier DROP (IF EXISTS)?
+        dropPartitionSpec (',' dropPartitionSpec)*                     #dropTablePartitions
     | ALTER TABLE tableIdentifier partitionSpec? SET locationSpec      #setTableLocation
     | ALTER TABLE tableIdentifier RECOVER PARTITIONS                   #recoverPartitions
     | DROP TABLE (IF EXISTS)? tableIdentifier PURGE?                   #dropTable
@@ -261,6 +261,14 @@ partitionVal
     : identifier (EQ constant)?
     ;
 
+dropPartitionSpec
+    : PARTITION '(' dropPartitionVal (',' dropPartitionVal)* ')'
+    ;
+
+dropPartitionVal
+    : identifier (comparisonOperator constant)?
+    ;
+
 describeFuncName
     : qualifiedName
     | STRING

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

@@ -293,6 +293,28 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
     }
   }
 
+  /**
+   * Create a partition specification map with filters.
+   */
+  override def visitDropPartitionSpec(
+      ctx: DropPartitionSpecContext): Seq[Expression] = {
+    withOrigin(ctx) {
+      ctx.dropPartitionVal().asScala.map { pFilter =>
+        if (pFilter.identifier() == null || pFilter.constant() == null ||
+            pFilter.comparisonOperator() == null) {
+          throw new ParseException(s"Invalid partition spec: ${pFilter.getText}", ctx)
+        }
+        // We cannot use UnresolvedAttribute because resolution is performed after Analysis, when
+        // running the command. The type is not relevant, it is replaced during the real resolution
+        val partition =
+          AttributeReference(pFilter.identifier().getText, StringType)()
+        val value = Literal(visitStringConstant(pFilter.constant()))
+        val operator = pFilter.comparisonOperator().getChild(0).asInstanceOf[TerminalNode]
+        buildComparison(partition, value, operator)
+      }
+    }
+  }
+
   /**
    * Convert a constant of any type into a string. This is typically used in DDL commands, and its
    * main purpose is to prevent slight differences due to back to back conversions i.e.:
@@ -1015,6 +1037,23 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
     val left = expression(ctx.left)
     val right = expression(ctx.right)
     val operator = ctx.comparisonOperator().getChild(0).asInstanceOf[TerminalNode]
+    buildComparison(left, right, operator)
+  }
+
+  /**
+   * Creates a comparison expression. The following comparison operators are supported:
+   * - Equal: '=' or '=='
+   * - Null-safe Equal: '<=>'
+   * - Not Equal: '<>' or '!='
+   * - Less than: '<'
+   * - Less then or Equal: '<='
+   * - Greater than: '>'
+   * - Greater then or Equal: '>='
+   */
+  private def buildComparison(
+      left: Expression,
+      right: Expression,
+      operator: TerminalNode): Expression = {
     operator.getSymbol.getType match {
       case SqlBaseParser.EQ =>
         EqualTo(left, right)

sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala

@@ -927,7 +927,7 @@ class SparkSqlAstBuilder(conf: SQLConf) extends AstBuilder(conf) {
     }
     AlterTableDropPartitionCommand(
       visitTableIdentifier(ctx.tableIdentifier),
-      ctx.partitionSpec.asScala.map(visitNonOptionalPartitionSpec),
+      ctx.dropPartitionSpec().asScala.map(visitDropPartitionSpec),
       ifExists = ctx.EXISTS != null,
       purge = ctx.PURGE != null,
       retainData = false)

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

@@ -29,10 +29,10 @@ import org.apache.hadoop.mapred.{FileInputFormat, JobConf}
 
 import org.apache.spark.sql.{AnalysisException, Row, SparkSession}
 import org.apache.spark.sql.catalyst.TableIdentifier
-import org.apache.spark.sql.catalyst.analysis.{NoSuchTableException, Resolver}
+import org.apache.spark.sql.catalyst.analysis.{Resolver, UnresolvedAttribute}
 import org.apache.spark.sql.catalyst.catalog._
 import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec
-import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation, PartitioningUtils}
 import org.apache.spark.sql.execution.datasources.orc.OrcFileFormat
@@ -521,35 +521,112 @@ case class AlterTableRenamePartitionCommand(
  */
 case class AlterTableDropPartitionCommand(
     tableName: TableIdentifier,
-    specs: Seq[TablePartitionSpec],
+    partitionsFilters: Seq[Seq[Expression]],
     ifExists: Boolean,
     purge: Boolean,
     retainData: Boolean)
   extends RunnableCommand {
 
   override def run(sparkSession: SparkSession): Seq[Row] = {
     val catalog = sparkSession.sessionState.catalog
+    val timeZone = Option(sparkSession.sessionState.conf.sessionLocalTimeZone)
     val table = catalog.getTableMetadata(tableName)
+    val partitionColumns = table.partitionColumnNames
+    val partitionAttributes = table.partitionSchema.toAttributes.map(a => a.name -> a).toMap
     DDLUtils.verifyAlterTableType(catalog, table, isView = false)
     DDLUtils.verifyPartitionProviderIsHive(sparkSession, table, "ALTER TABLE DROP PARTITION")
 
-    val normalizedSpecs = specs.map { spec =>
-      PartitioningUtils.normalizePartitionSpec(
-        spec,
-        table.partitionColumnNames,
-        table.identifier.quotedString,
-        sparkSession.sessionState.conf.resolver)
+    val resolvedSpecs = partitionsFilters.flatMap { filtersSpec =>
+      if (hasComplexFilters(filtersSpec)) {
+        generatePartitionSpec(filtersSpec,
+          partitionColumns,
+          partitionAttributes,
+          table.identifier,
+          catalog,
+          sparkSession.sessionState.conf.resolver,
+          timeZone,
+          ifExists)
+      } else {
+        val partitionSpec = filtersSpec.map {
+          case EqualTo(key: Attribute, Literal(value, StringType)) =>
+            key.name -> value.toString
+        }.toMap
+        PartitioningUtils.normalizePartitionSpec(
+          partitionSpec,
+          partitionColumns,
+          table.identifier.quotedString,
+          sparkSession.sessionState.conf.resolver) :: Nil
+      }
     }
 
     catalog.dropPartitions(
-      table.identifier, normalizedSpecs, ignoreIfNotExists = ifExists, purge = purge,
+      table.identifier, resolvedSpecs, ignoreIfNotExists = ifExists, purge = purge,
       retainData = retainData)
 
     CommandUtils.updateTableStats(sparkSession, table)
 
     Seq.empty[Row]
   }
 
+  def hasComplexFilters(partitionFilterSpec: Seq[Expression]): Boolean = {
+    partitionFilterSpec.exists(!_.isInstanceOf[EqualTo])
+  }
+
+  def generatePartitionSpec(
+      partitionFilterSpec: Seq[Expression],
+      partitionColumns: Seq[String],
+      partitionAttributes: Map[String, Attribute],
+      tableIdentifier: TableIdentifier,
+      catalog: SessionCatalog,
+      resolver: Resolver,
+      timeZone: Option[String],
+      ifExists: Boolean): Seq[TablePartitionSpec] = {
+    val filters = partitionFilterSpec.map { pFilter =>
+      pFilter.transform {
+        // Resolve the partition attributes
+        case partitionCol: Attribute =>
+          val normalizedPartition = PartitioningUtils.normalizePartitionColumn(
+            partitionCol.name,
+            partitionColumns,
+            tableIdentifier.quotedString,
+            resolver)
+          partitionAttributes(normalizedPartition)
+      }.transform {
+        // Cast the partition value to the data type of the corresponding partition attribute
+        case cmp @ BinaryComparison(partitionAttr, value)
+            if !partitionAttr.dataType.sameType(value.dataType) =>
+          cmp.withNewChildren(Seq(partitionAttr, Cast(value, partitionAttr.dataType, timeZone)))
+      }
+    }
+    val partitions = catalog.listPartitionsByFilter(tableIdentifier, filters)
+    if (partitions.isEmpty && !ifExists) {
+      throw new AnalysisException(s"There is no partition for ${filters.reduceLeft(And).sql}")
+    }
+    partitions.map(_.spec)
+  }
+}
+
+
+object AlterTableDropPartitionCommand {
+
+  def fromSpecs(
+      tableName: TableIdentifier,
+      specs: Seq[TablePartitionSpec],
+      ifExists: Boolean,
+      purge: Boolean,
+      retainData: Boolean): AlterTableDropPartitionCommand = {
+    AlterTableDropPartitionCommand(tableName,
+      specs.map(tablePartitionToPartitionFilters),
+      ifExists,
+      purge,
+      retainData)
+  }
+
+  def tablePartitionToPartitionFilters(spec: TablePartitionSpec): Seq[Expression] = {
+    spec.map {
+      case (key, value) => EqualTo(AttributeReference(key, StringType)(), Literal(value))
+    }.toSeq
+  }
 }
 
 

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

@@ -141,7 +141,7 @@ case class InsertIntoHadoopFsRelationCommand(
           if (mode == SaveMode.Overwrite && !dynamicPartitionOverwrite) {
             val deletedPartitions = initialMatchingPartitions.toSet -- updatedPartitions
             if (deletedPartitions.nonEmpty) {
-              AlterTableDropPartitionCommand(
+              AlterTableDropPartitionCommand.fromSpecs(
                 catalogTable.get.identifier, deletedPartitions.toSeq,
                 ifExists = true, purge = false,
                 retainData = true /* already deleted */).run(sparkSession)

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

@@ -296,9 +296,7 @@ object PartitioningUtils {
       tblName: String,
       resolver: Resolver): Map[String, T] = {
     val normalizedPartSpec = partitionSpec.toSeq.map { case (key, value) =>
-      val normalizedKey = partColNames.find(resolver(_, key)).getOrElse {
-        throw new AnalysisException(s"$key is not a valid partition column in table $tblName.")
-      }
+      val normalizedKey = normalizePartitionColumn(key, partColNames, tblName, resolver)
       normalizedKey -> value
     }
 
@@ -308,6 +306,16 @@ object PartitioningUtils {
     normalizedPartSpec.toMap
   }
 
+  def normalizePartitionColumn(
+      partition: String,
+      partColNames: Seq[String],
+      tblName: String,
+      resolver: Resolver): String = {
+    partColNames.find(resolver(_, partition)).getOrElse {
+      throw new AnalysisException(s"$partition is not a valid partition column in table $tblName.")
+    }
+  }
+
   /**
    * Resolves possible type conflicts between partitions by up-casting "lower" types using
    * [[findWiderTypeForPartitionColumn]].

sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLParserSuite.scala

@@ -861,7 +861,8 @@ class DDLParserSuite extends PlanTest with SharedSQLContext {
     assertUnsupported(sql2_view)
 
     val tableIdent = TableIdentifier("table_name", None)
-    val expected1_table = AlterTableDropPartitionCommand(
+
+    val expected1_table = AlterTableDropPartitionCommand.fromSpecs(
       tableIdent,
       Seq(
         Map("dt" -> "2008-08-08", "country" -> "us"),