From 0e8c9316533db74370e2f49ac920a1e1e1c3a1e7 Mon Sep 17 00:00:00 2001
From: Sabir <sabir.akhadov@databricks.com>
Date: Wed, 28 Feb 2024 12:37:48 +0100
Subject: [PATCH] .

---
 .../commands/alterDeltaTableCommands.scala    |  4 ++--
 .../RemoveColumnMappingCommand.scala          | 19 ++++++++++++++++---
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/commands/alterDeltaTableCommands.scala b/spark/src/main/scala/org/apache/spark/sql/delta/commands/alterDeltaTableCommands.scala
index 097b36b4404..028e7be9eb9 100644
--- a/spark/src/main/scala/org/apache/spark/sql/delta/commands/alterDeltaTableCommands.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/delta/commands/alterDeltaTableCommands.scala
@@ -117,7 +117,7 @@ case class AlterTableSetPropertiesDeltaCommand(
     val columnMappingRemovalAllowed = sparkSession.sessionState.conf.getConf(
       DeltaSQLConf.ALLOW_COLUMN_MAPPING_REMOVAL)
     if (disableColumnMapping && columnMappingRemovalAllowed) {
-      new RemoveColumnMappingCommand(deltaLog, table.catalogTable)
+      RemoveColumnMappingCommand(deltaLog, table.catalogTable)
         .run(sparkSession, removeColumnMappingTableProperty = false)
     }
     recordDeltaOperation(deltaLog, "delta.ddl.alter.setProperties") {
@@ -176,7 +176,7 @@ case class AlterTableUnsetPropertiesDeltaCommand(
     val columnMappingRemovalAllowed = sparkSession.sessionState.conf.getConf(
       DeltaSQLConf.ALLOW_COLUMN_MAPPING_REMOVAL)
     if (disableColumnMapping && columnMappingRemovalAllowed) {
-      new RemoveColumnMappingCommand(deltaLog, table.catalogTable)
+      RemoveColumnMappingCommand(deltaLog, table.catalogTable)
         .run(sparkSession, removeColumnMappingTableProperty = true)
     }
     recordDeltaOperation(deltaLog, "delta.ddl.alter.unsetProperties") {
diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/commands/columnmapping/RemoveColumnMappingCommand.scala b/spark/src/main/scala/org/apache/spark/sql/delta/commands/columnmapping/RemoveColumnMappingCommand.scala
index da5c95ed5b4..d989f9cffde 100644
--- a/spark/src/main/scala/org/apache/spark/sql/delta/commands/columnmapping/RemoveColumnMappingCommand.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/delta/commands/columnmapping/RemoveColumnMappingCommand.scala
@@ -19,11 +19,16 @@ package org.apache.spark.sql.delta.commands.columnmapping
 import org.apache.spark.sql.delta.{DeltaErrors, DeltaLog}
 import org.apache.spark.sql.delta.schema.{ImplicitMetadataOperation, SchemaUtils}
 
-import org.apache.spark.sql.{Row, SparkSession}
+import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.catalog.CatalogTable
 import org.apache.spark.sql.types.StructType
 
-class RemoveColumnMappingCommand(deltaLog: DeltaLog, catalogOpt: Option[CatalogTable])
+/**
+ * A command to remove the column mapping from a table.
+ */
+class RemoveColumnMappingCommand(
+    val deltaLog: DeltaLog,
+    val catalogOpt: Option[CatalogTable])
   extends ImplicitMetadataOperation {
   override protected val canMergeSchema: Boolean = false
   override protected val canOverwriteSchema: Boolean = true
@@ -41,7 +46,7 @@ class RemoveColumnMappingCommand(deltaLog: DeltaLog, catalogOpt: Option[CatalogT
   /**
    * Verify none of the schema fields contain invalid column names.
    */
-  private def verifySchemaFieldNames(schema: StructType) = {
+  protected def verifySchemaFieldNames(schema: StructType) = {
     val invalidColumnNames =
       SchemaUtils.findInvalidColumnNamesInSchema(schema)
     if (invalidColumnNames.nonEmpty) {
@@ -50,3 +55,11 @@ class RemoveColumnMappingCommand(deltaLog: DeltaLog, catalogOpt: Option[CatalogT
     }
   }
 }
+
+object RemoveColumnMappingCommand {
+  def apply(
+      deltaLog: DeltaLog,
+      catalogOpt: Option[CatalogTable]): RemoveColumnMappingCommand = {
+    new RemoveColumnMappingCommand(deltaLog, catalogOpt)
+  }
+}