apache · mgaido91 · Apr 5, 2018 · Sep 5, 2018 · Sep 6, 2018 · Sep 9, 2018
diff --git a/sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4 b/sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4
@@ -149,10 +149,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 multipartIdentifier SET locationSpec                 #setTableLocation
     | ALTER TABLE tableIdentifier partitionSpec SET locationSpec       #setPartitionLocation
     | ALTER TABLE tableIdentifier RECOVER PARTITIONS                   #recoverPartitions
@@ -300,6 +300,14 @@ partitionVal
     : identifier (EQ constant)?
     ;
 
+dropPartitionSpec
+    : PARTITION '(' dropPartitionVal (',' dropPartitionVal)* ')'
+    ;
+
+dropPartitionVal
+    : identifier (comparisonOperator constant)?
+    ;
+
 database
     : DATABASE
     | SCHEMA

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/namedExpressions.scala
@@ -391,6 +391,35 @@ case class OuterReference(e: NamedExpression)
   override def newInstance(): NamedExpression = OuterReference(e.newInstance())
 }
 
+/**
+ * A place holder used to hold the name of the partition attributes specified when running commands
+ * involving partitions, eg. ALTER TABLE ... DROP PARTITIONS.
+ */
+case class PartitioningAttribute(
+    name: String,
+    override val exprId: ExprId = NamedExpression.newExprId)
+  extends Attribute with Unevaluable {
+  // We need a dataType to be used during analysis for resolving the expressions (see
+  // checkInputDataTypes). The String type is used because all the literals in PARTITION operations
+  // are parsed as strings and eventually casted later.
+  override def dataType: DataType = StringType
+  override def nullable: Boolean = false
+
+  override def qualifier: Seq[String] = throw new UnsupportedOperationException
+  override def withNullability(newNullability: Boolean): Attribute =
+    throw new UnsupportedOperationException
+  override def newInstance(): Attribute = throw new UnsupportedOperationException
+  override def withQualifier(newQualifier: Seq[String]): Attribute =
+    throw new UnsupportedOperationException
+  override def withName(newName: String): Attribute = throw new UnsupportedOperationException
+  override def withMetadata(newMetadata: Metadata): Attribute =
+    throw new UnsupportedOperationException
+
+  override lazy val canonicalized: Expression = this.copy(exprId = ExprId(0))
+
+  override def withExprId(newExprId: ExprId): Attribute = throw new UnsupportedOperationException
+}
+
 object VirtualColumn {
   // The attribute name used by Hive, which has different result than Spark, deprecated.
   val hiveGroupingIdName: String = "grouping__id"

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala
@@ -352,6 +352,29 @@ 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.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.
+      val partition = PartitioningAttribute(pFilter.identifier().getText)
+      val value = Literal(visitStringConstant(pFilter.constant()))
+      val operator = pFilter.comparisonOperator().getChild(0).asInstanceOf[TerminalNode]
+      val comparison = buildComparison(partition, value, operator)
+      if (comparison.isInstanceOf[EqualNullSafe]) {
+        throw new ParseException(
+          "'<=>' operator is not supported in ALTER TABLE ... DROP PARTITION.", ctx)
+      }
+      comparison
+    }
+  }
+
   /**
    * 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.:
@@ -1174,6 +1197,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)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkSqlParser.scala
@@ -749,7 +749,7 @@ class SparkSqlAstBuilder(conf: SQLConf) extends AstBuilder(conf) {
     }
     AlterTableDropPartitionCommand(
       visitTableIdentifier(ctx.tableIdentifier),
-      ctx.partitionSpec.asScala.map(visitNonOptionalPartitionSpec),
+      ctx.dropPartitionSpec().asScala.map(visitDropPartitionSpec),
    *   ALTER TABLE table DROP [IF EXISTS] PARTITION spec1[, PARTITION spec2, ...] [PURGE]; 
    *   ALTER TABLE table DROP [IF EXISTS] PARTITION spec1[, PARTITION spec2, ...] [PURGE]; 
       ifExists = ctx.EXISTS != null,
       purge = ctx.PURGE != null,
       retainData = false)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ddl.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/command/ddl.scala
@@ -31,10 +31,10 @@ import org.apache.hadoop.mapred.{FileInputFormat, JobConf}
 import org.apache.spark.internal.config.RDD_PARALLEL_LISTING_THRESHOLD
 import org.apache.spark.sql.{AnalysisException, Row, SparkSession}
 import org.apache.spark.sql.catalyst.TableIdentifier
-import org.apache.spark.sql.catalyst.analysis.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
@@ -524,35 +524,114 @@ 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: PartitioningAttribute =>
+          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) =>
+          val dt = partitionAttr.dataType
+          val lit = Literal(Cast(value, dt, timeZone).eval(), dt)
+          cmp.withNewChildren(Seq(partitionAttr, lit))
+      }
+    }
+    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(PartitioningAttribute(key), Literal(value))
+    }.toSeq
+  }
 }
 
 

diff --git a/.../scala/org/apache/spark/sql/execution/datasources/InsertIntoHadoopFsRelationCommand.scala b/.../scala/org/apache/spark/sql/execution/datasources/InsertIntoHadoopFsRelationCommand.scala
@@ -148,7 +148,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)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/PartitioningUtils.scala
@@ -366,9 +366,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
     }
 
@@ -378,6 +376,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]].

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLParserSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/DDLParserSuite.scala
@@ -29,7 +29,7 @@ import org.apache.spark.sql.catalyst.catalog._
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans
 import org.apache.spark.sql.catalyst.dsl.plans.DslLogicalPlan
-import org.apache.spark.sql.catalyst.expressions.JsonTuple
+import org.apache.spark.sql.catalyst.expressions.{Expression, JsonTuple, PartitioningAttribute}
 import org.apache.spark.sql.catalyst.parser.ParseException
 import org.apache.spark.sql.catalyst.plans.logical.{Generate, InsertIntoDir, LogicalPlan, Project, ScriptTransformation}
 import org.apache.spark.sql.execution.SparkSqlParser
@@ -687,7 +687,8 @@ class DDLParserSuite extends AnalysisTest 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"),
@@ -698,9 +699,36 @@ class DDLParserSuite extends AnalysisTest with SharedSQLContext {
     val expected2_table = expected1_table.copy(ifExists = false)
     val expected1_purge = expected1_table.copy(purge = true)
 
-    comparePlans(parsed1_table, expected1_table)
-    comparePlans(parsed2_table, expected2_table)
-    comparePlans(parsed1_purge, expected1_purge)
+    comparePlans(parsed1_table.canonicalized, expected1_table.canonicalized)
+    comparePlans(parsed2_table.canonicalized, expected2_table.canonicalized)
+    comparePlans(parsed1_purge.canonicalized, expected1_purge.canonicalized)
+
+    // SPARK-23866: Support any comparison operator in ALTER TABLE ... DROP PARTITION
+    Seq((">", (a: Expression, b: Expression) => a > b),
+      (">=", (a: Expression, b: Expression) => a >= b),
+      ("<", (a: Expression, b: Expression) => a < b),
+      ("<=", (a: Expression, b: Expression) => a <= b),
+      ("<>", (a: Expression, b: Expression) => a =!= b),
+      ("!=", (a: Expression, b: Expression) => a =!= b)).foreach { case (op, predicateGen) =>
+        val genPlan = parser.parsePlan(sql1_table.replace("=", op))
+        val dtAttr = PartitioningAttribute("dt")
+        val countryAttr = PartitioningAttribute("country")
+        val expectedPlan = AlterTableDropPartitionCommand(
+          tableIdent,
+          Seq(
+            Seq(predicateGen(dtAttr, "2008-08-08"), predicateGen(countryAttr, "us")),
+            Seq(predicateGen(dtAttr, "2009-09-09"), predicateGen(countryAttr, "uk"))),
+          ifExists = true,
+          purge = false,
+          retainData = false)
+        comparePlans(genPlan.canonicalized, expectedPlan.canonicalized)
+    }
+
+    // SPARK-23866: <=> is not supported
+    intercept("ALTER TABLE table_name DROP PARTITION (dt <=> 'a')", "operator is not supported in")
+
+    // SPARK-23866: Invalid partition specification
+    intercept("ALTER TABLE table_name DROP PARTITION (dt)", "Invalid partition spec:")
   }
 
   test("alter table: archive partition (not supported)") {