[spark] Support drop partition from top level

Author: zhangyongxiang.alpha

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/analysis/PaimonAnalysis.scala

@@ -21,6 +21,7 @@ package org.apache.paimon.spark.catalyst.analysis
 import org.apache.paimon.spark.{SparkConnectorOptions, SparkTable}
 import org.apache.paimon.spark.catalyst.Compatibility
 import org.apache.paimon.spark.catalyst.analysis.PaimonRelation.isPaimonTable
+import org.apache.paimon.spark.catalyst.plans.logical.PaimonDropPartitions
 import org.apache.paimon.spark.commands.{PaimonAnalyzeTableColumnCommand, PaimonDynamicPartitionOverwriteCommand, PaimonShowColumnsCommand, PaimonTruncateTableCommand}
 import org.apache.paimon.spark.util.OptionUtils
 import org.apache.paimon.table.FileStoreTable
@@ -33,13 +34,13 @@ import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.CharVarcharUtils
 import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLId
 import org.apache.spark.sql.connector.catalog.TableCapability
-import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits, DataSourceV2Relation}
 import org.apache.spark.sql.types._
 
 import scala.collection.mutable
 
 class PaimonAnalysis(session: SparkSession) extends Rule[LogicalPlan] {
-
+  import DataSourceV2Implicits._
   override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsDown {
 
     case a @ PaimonV2WriteCommand(table) if !paimonWriteResolved(a.query, table) =>
@@ -59,6 +60,11 @@ class PaimonAnalysis(session: SparkSession) extends Rule[LogicalPlan] {
 
     case s @ ShowColumns(PaimonRelation(table), _, _) if s.resolved =>
       PaimonShowColumnsCommand(table)
+
+    case d @ PaimonDropPartitions(ResolvedTable(_, _, table: SparkTable, _), parts, _, _)
+        if d.resolved =>
+      PaimonDropPartitions.validate(table, parts.asResolvedPartitionSpecs)
+      d
   }
 
   private def writeOptions(v2WriteCommand: V2WriteCommand): Map[String, String] = {

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/catalyst/plans/logical/PaimonDropPartitions.scala

@@ -0,0 +1,59 @@
+/*
+ * 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.paimon.spark.catalyst.plans.logical
+
+import org.apache.paimon.spark.SparkTable
+
+import org.apache.spark.sql.{types, PaimonUtils}
+import org.apache.spark.sql.catalyst.analysis.{PartitionSpec, ResolvedPartitionSpec}
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, V2PartitionCommand}
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits.TableHelper
+import org.apache.spark.sql.types.StructType;
+
+case class PaimonDropPartitions(
+    table: LogicalPlan,
+    parts: Seq[PartitionSpec],
+    ifExists: Boolean,
+    purge: Boolean)
+  extends V2PartitionCommand {
+
+  override def allowPartialPartitionSpec: Boolean = true
+  override protected def withNewChildInternal(newChild: LogicalPlan): PaimonDropPartitions =
+    copy(table = newChild)
+}
+
+object PaimonDropPartitions {
+  def validate(table: SparkTable, partialSpecs: Seq[ResolvedPartitionSpec]): Unit = {
+    val partitionSchema = table.asPartitionable.partitionSchema();
+    partialSpecs.foreach {
+      partialSpec =>
+        if (!partitionSchema.names.toSeq.startsWith(partialSpec.names)) {
+          val values = partialSpec.names.zipWithIndex.map {
+            case (name, ordinal) =>
+              partialSpec.ident.get(ordinal, partitionSchema.apply(name).dataType).toString
+          }
+          val spec = partialSpec.names
+            .zip(values)
+            .map { case (name, value) => s"$name = '$value'" }
+            .mkString(",")
+          throw PaimonUtils.invalidPartitionSpecError(spec, partitionSchema.fieldNames, table.name)
+        }
+    }
+  }
+}

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonDropPartitionsExec.scala

@@ -0,0 +1,84 @@
+/*
+ * 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.paimon.spark.execution
+
+import org.apache.paimon.spark.SparkTable
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.analysis.{NoSuchPartitionsException, ResolvedPartitionSpec}
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits.TableHelper
+import org.apache.spark.sql.execution.datasources.v2.LeafV2CommandExec
+
+case class PaimonDropPartitionsExec(
+    table: SparkTable,
+    partSpecs: Seq[ResolvedPartitionSpec],
+    ignoreIfNotExists: Boolean,
+    purge: Boolean,
+    refreshCache: () => Unit)
+  extends LeafV2CommandExec
+  with Logging {
+  override protected def run(): Seq[InternalRow] = {
+    val partitionSchema = table.asPartitionable.partitionSchema()
+    val (partialPartSpecs, fullPartSpecs) =
+      partSpecs.partition(_.ident.numFields != partitionSchema.length)
+
+    val (existsPartIdents, notExistsPartIdents) =
+      fullPartSpecs.map(_.ident).partition(table.partitionExists)
+    if (notExistsPartIdents.nonEmpty && !ignoreIfNotExists) {
+      throw new NoSuchPartitionsException(
+        table.name(),
+        notExistsPartIdents,
+        table.asPartitionable.partitionSchema())
+    }
+    val allExistsPartIdents = existsPartIdents ++ partialPartSpecs.flatMap(expendPartialSpec)
+    logInfo("Try to drop partitions: " + allExistsPartIdents.mkString(","))
+    val isTableAltered = if (allExistsPartIdents.nonEmpty) {
+      allExistsPartIdents
+        .map(
+          partIdent => {
+            if (purge) table.purgePartition(partIdent) else table.dropPartition(partIdent)
+          })
+        .reduce(_ || _)
+    } else false
+
+    if (isTableAltered) refreshCache()
+    Seq.empty
+  }
+
+  private def expendPartialSpec(partialSpec: ResolvedPartitionSpec): Seq[InternalRow] = {
+    assertSpec(partialSpec)
+    table.listPartitionIdentifiers(partialSpec.names.toArray, partialSpec.ident).toSeq
+  }
+
+  private def assertSpec(partialSpec: ResolvedPartitionSpec): Unit = {
+    val partitionSchema = table.asPartitionable.partitionSchema();
+    if (partialSpec.ident.numFields > partitionSchema.length) {
+      throw new IllegalArgumentException(
+        s"Partial partition spec should be part of $partitionSchema fields, but got $partialSpec")
+    }
+    if (!partitionSchema.names.toSeq.startsWith(partialSpec.names)) {
+      throw new IllegalArgumentException(
+        s"Partial partition spec should be prefix of $partitionSchema fields, but got $partialSpec")
+    }
+  }
+
+  override def output: Seq[Attribute] = Seq.empty
+}

paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/execution/PaimonStrategy.scala

@@ -18,19 +18,22 @@
 
 package org.apache.paimon.spark.execution
 
-import org.apache.paimon.spark.{SparkCatalog, SparkGenericCatalog, SparkUtils}
+import org.apache.paimon.spark.{SparkCatalog, SparkGenericCatalog, SparkTable, SparkUtils}
 import org.apache.paimon.spark.catalog.SupportView
 import org.apache.paimon.spark.catalyst.analysis.ResolvedPaimonView
 import org.apache.paimon.spark.catalyst.plans.logical.{CreateOrReplaceTagCommand, CreatePaimonView, DeleteTagCommand, DropPaimonView, PaimonCallCommand, RenameTagCommand, ResolvedIdentifier, ShowPaimonViews, ShowTagsCommand}
+import org.apache.paimon.spark.catalyst.plans.logical.PaimonDropPartitions
 
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.analysis.ResolvedNamespace
+import org.apache.spark.sql.catalyst.analysis.{ResolvedNamespace, ResolvedTable}
 import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow, PredicateHelper}
 import org.apache.spark.sql.catalyst.plans.logical.{CreateTableAsSelect, DescribeRelation, LogicalPlan, ShowCreateTable}
 import org.apache.spark.sql.connector.catalog.{Identifier, PaimonLookupCatalog, TableCatalog}
 import org.apache.spark.sql.execution.{SparkPlan, SparkStrategy}
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits, DataSourceV2Relation}
 import org.apache.spark.sql.execution.shim.PaimonCreateTableAsSelectStrategy
+import org.apache.spark.sql.paimon.shims.SparkShimLoader
 
 import scala.collection.JavaConverters._
 
@@ -39,6 +42,7 @@ case class PaimonStrategy(spark: SparkSession)
   with PredicateHelper
   with PaimonLookupCatalog {
 
+  import DataSourceV2Implicits._
   protected lazy val catalogManager = spark.sessionState.catalogManager
 
   override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
@@ -107,6 +111,18 @@ case class PaimonStrategy(spark: SparkSession)
     case DescribeRelation(ResolvedPaimonView(viewCatalog, ident), _, isExtended, output) =>
       DescribePaimonViewExec(output, viewCatalog, ident, isExtended) :: Nil
 
+    case d @ PaimonDropPartitions(
+          r @ ResolvedTable(_, _, table: SparkTable, _),
+          parts,
+          ifExists,
+          purge) =>
+      PaimonDropPartitionsExec(
+        table,
+        parts.asResolvedPartitionSpecs,
+        ifExists,
+        purge,
+        recacheTable(r)) :: Nil
+
     case _ => Nil
   }
 
@@ -132,4 +148,9 @@ case class PaimonStrategy(spark: SparkSession)
       }
     }
   }
+
+  private def recacheTable(r: ResolvedTable)(): Unit = {
+    val v2Relation = DataSourceV2Relation.create(r.table, Some(r.catalog), Some(r.identifier))
+    SparkShimLoader.shim.classicApi.recacheByPlan(spark, v2Relation)
+  }
 }

paimon-spark/paimon-spark-common/src/main/scala/org/apache/spark/sql/PaimonUtils.scala

@@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression}
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.connector.expressions.FieldReference
 import org.apache.spark.sql.connector.expressions.filter.Predicate
+import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.execution.datasources.DataSourceStrategy
 import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Strategy.translateFilterV2WithMapping
 import org.apache.spark.sql.internal.connector.PredicateUtils
@@ -134,4 +135,11 @@ object PaimonUtils {
   def classForName(clazz: String): Class[_] = {
     SparkUtils.classForName(clazz)
   }
+
+  def invalidPartitionSpecError(
+      specKeys: String,
+      partitionColumnNames: Seq[String],
+      tableName: String): Throwable = {
+    QueryCompilationErrors.invalidPartitionSpecError(specKeys, partitionColumnNames, tableName)
+  }
 }

paimon-spark/paimon-spark-common/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/AbstractPaimonSparkSqlExtensionsParser.scala

@@ -79,7 +79,8 @@ abstract class AbstractPaimonSparkSqlExtensionsParser(val delegate: ParserInterf
   private def parserRules(sparkSession: SparkSession): Seq[Rule[LogicalPlan]] = {
     Seq(
       RewritePaimonViewCommands(sparkSession),
-      RewritePaimonFunctionCommands(sparkSession)
+      RewritePaimonFunctionCommands(sparkSession),
+      RewriteSparkDDLCommands(sparkSession)
     )
   }
 

paimon-spark/paimon-spark-common/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/RewriteSparkDDLCommands.scala

@@ -0,0 +1,44 @@
+/*
+ * 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.catalyst.parser.extensions
+
+import org.apache.paimon.spark.catalog.SupportView
+import org.apache.paimon.spark.catalyst.plans.logical.{PaimonDropPartitions, ResolvedIdentifier}
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.analysis.UnresolvedIdentifier
+import org.apache.spark.sql.catalyst.plans.logical.{DropPartitions, LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.connector.catalog.{CatalogManager, LookupCatalog}
+
+case class RewriteSparkDDLCommands(spark: SparkSession)
+  extends Rule[LogicalPlan]
+  with LookupCatalog {
+
+  protected lazy val catalogManager: CatalogManager = spark.sessionState.catalogManager
+
+  override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsUp {
+
+    // A new member was added to CreatePaimonView since spark4.0,
+    // unapply pattern matching is not used here to ensure compatibility across multiple spark versions.
+    case DropPartitions(UnresolvedPaimonRelation(aliasedTable), parts, ifExists, purge) =>
+      PaimonDropPartitions(aliasedTable, parts, ifExists, purge)
+
+  }
+}

paimon-spark/paimon-spark-common/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/UnresolvedPaimonRelation.scala

@@ -0,0 +1,49 @@
+/*
+ * 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.catalyst.parser.extensions
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.analysis.{EliminateSubqueryAliases, UnresolvedTable}
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.connector.catalog.{PaimonLookupCatalog, TableCatalog}
+
+import scala.util.Try
+
+object UnresolvedPaimonRelation extends PaimonLookupCatalog {
+  protected lazy val catalogManager = SparkSession.active.sessionState.catalogManager
+
+  def unapply(plan: LogicalPlan): Option[LogicalPlan] = {
+    EliminateSubqueryAliases(plan) match {
+      case UnresolvedTable(multipartIdentifier, _, _) if isPaimonTable(multipartIdentifier) =>
+        Some(plan)
+      case _ => None
+    }
+  }
+
+  private def isPaimonTable(multipartIdentifier: Seq[String]): Boolean = {
+    multipartIdentifier match {
+      case CatalogAndIdentifier(catalog: TableCatalog, ident) =>
+        Try(catalog.loadTable(ident))
+          .map(t => t.isInstanceOf[org.apache.paimon.spark.SparkTable])
+          .getOrElse(false)
+      case _ => false
+    }
+  }
+
+}