[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/catalog/interface.scala

@@ -552,6 +552,11 @@ object CatalogTypes {
    */
   type TablePartitionSpec = Map[String, String]
 
+  /**
+   * Specifications of table partition filters. Seq of column name, comparison operator and value.
+   */
+  type PartitionFiltersSpec = Seq[(String, String, String)]
+
   /**
    * Initialize an empty spec.
    */

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

@@ -293,6 +293,22 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
     }
   }
 
+  /**
+   * Create a partition specification map with filters.
+   */
+  override def visitDropPartitionSpec(
+      ctx: DropPartitionSpecContext): Seq[(String, String, String)] = {
+    withOrigin(ctx) {
+      ctx.dropPartitionVal().asScala.map { pFilter =>
+        val partition = pFilter.identifier().getText
+        val value = visitStringConstant(pFilter.constant())
+        val operator = pFilter.comparisonOperator().getChild(0).asInstanceOf[TerminalNode]
+        val stringOperator = SqlBaseParser.VOCABULARY.getSymbolicName(operator.getSymbol.getType)
+        (partition, stringOperator, value)
+      }
+    }
+  }
+
   /**
    * 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.:

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
 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.catalog.CatalogTypes.{PartitionFiltersSpec, TablePartitionSpec}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Cast, EqualNullSafe, EqualTo, GreaterThan, GreaterThanOrEqual, LessThan, LessThanOrEqual, Literal, Not}
 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,114 @@ case class AlterTableRenamePartitionCommand(
  */
 case class AlterTableDropPartitionCommand(
     tableName: TableIdentifier,
-    specs: Seq[TablePartitionSpec],
+    partitionsFilters: Seq[PartitionFiltersSpec],
     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)
+      } else {
+        val partitionSpec = filtersSpec.map {
+          case (key, _, value) => key -> value
+        }.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: PartitionFiltersSpec): Boolean = {
+    !partitionFilterSpec.forall(_._2 == "EQ")
+  }
+
+  def generatePartitionSpec(
+      partitionFilterSpec: PartitionFiltersSpec,
+      partitionColumns: Seq[String],
+      partitionAttributes: Map[String, Attribute],
+      tableIdentifier: TableIdentifier,
+      catalog: SessionCatalog,
+      resolver: Resolver,
+      timeZone: Option[String]): Seq[TablePartitionSpec] = {
+    val filters = partitionFilterSpec.map { case (partitionColumn, operator, value) =>
+      val normalizedPartition = PartitioningUtils.normalizePartitionColumn(
+        partitionColumn,
+        partitionColumns,
+        tableIdentifier.quotedString,
+        resolver)
+      val partitionAttr = partitionAttributes(normalizedPartition)
+      val castedLiteralValue = Cast(Literal(value), partitionAttr.dataType, timeZone)
+      operator match {
+        case "EQ" =>
+          EqualTo(partitionAttr, castedLiteralValue)
+        case "NSEQ" =>
+          EqualNullSafe(partitionAttr, castedLiteralValue)
+        case "NEQ" | "NEQJ" =>
+          Not(EqualTo(partitionAttr, castedLiteralValue))
+        case "LT" =>
+          LessThan(partitionAttr, castedLiteralValue)
+        case "LTE" =>
+          LessThanOrEqual(partitionAttr, castedLiteralValue)
+        case "GT" =>
+          GreaterThan(partitionAttr, castedLiteralValue)
+        case "GTE" =>
+          GreaterThanOrEqual(partitionAttr, castedLiteralValue)
+      }
+    }
+    val partitions = catalog.listPartitionsByFilter(tableIdentifier, filters)
+    partitions.map(_.spec)
+  }
+}
+
+
+object AlterTableDropPartitionCommand {
+
+  def fromSpecs(
+      tableName: TableIdentifier,
+      specs: Seq[TablePartitionSpec],
+      ifExists: Boolean,
+      purge: Boolean,
+      retainData: Boolean): AlterTableDropPartitionCommand = {
+    AlterTableDropPartitionCommand(tableName,
+      specs.map(tablePartitionToPartitionFiltersSpec),
+      ifExists,
+      purge,
+      retainData)
+  }
+
+  def tablePartitionToPartitionFiltersSpec(spec: TablePartitionSpec): PartitionFiltersSpec = {
+    spec.map {
+      case (key, value) => (key, "EQ", 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"),