initial commit

Author: zedtang

spark/src/main/scala/org/apache/spark/sql/delta/DomainMetadataUtils.scala

@@ -16,6 +16,7 @@
 
 package org.apache.spark.sql.delta
 
+import org.apache.spark.sql.delta.skipping.clustering.ClusteredTableUtils
 import org.apache.spark.sql.delta.actions.{Action, DomainMetadata, Protocol}
 import org.apache.spark.sql.delta.clustering.ClusteringMetadataDomain
 import org.apache.spark.sql.delta.metering.DeltaLogging
@@ -91,15 +92,40 @@ object DomainMetadataUtils extends DeltaLogging {
 
   /**
    * Generates a new sequence of DomainMetadata to commits for RESTORE TABLE.
-   *  - Source (table to restore to) domains will be copied if they appear in the pre-defined
+   *  - Domains in the toSnapshot will be copied if they appear in the pre-defined
    *    "copy" list (e.g., table features require some specific domains to be copied).
-   *  - All other domains not in the list are "retained".
+   *  - All other domains not in the list are dropped from the "toSnapshot".
+   *
+   * For clustering metadata domains, it overwrites the existing domain metadata in the
+   * fromSnapshot with the following clustering columns.
+   * 1. If toSnapshot is not a clustered table or missing domain metadata, use empty clustering
+   *    columns.
+   * 2. If toSnapshot is a clustered table, use the clustering columns from toSnapshot.
+   *
+   * @param toSnapshot    The snapshot being restored to, which is referred as "source" table.
+   * @param fromSnapshot  The snapshot being restored from, which is the current state.
    */
   def handleDomainMetadataForRestoreTable(
-      sourceDomainMetadatas: Seq[DomainMetadata]): Seq[DomainMetadata] = {
-    sourceDomainMetadatas.filter { m =>
+      toSnapshot: Snapshot,
+      fromSnapshot: Snapshot): Seq[DomainMetadata] = {
+    val filteredDomainMetadata = toSnapshot.domainMetadata.filter { m =>
       METADATA_DOMAIN_TO_COPY_FOR_RESTORE_TABLE.contains(m.domain)
     }
+    val clusteringColumnsToRestore = ClusteredTableUtils.getClusteringColumnsOptional(toSnapshot)
+
+    val isRestoringToClusteredTable =
+      ClusteredTableUtils.isSupported(toSnapshot.protocol) && clusteringColumnsToRestore.nonEmpty
+    val clusteringColumns = if (isRestoringToClusteredTable) {
+      // We overwrite the clustering columns in the fromSnapshot with the clustering columns
+      // in the toSnapshot.
+      clusteringColumnsToRestore.get
+    } else {
+      // toSnapshot is not a clustered table or missing domain metadata, so we write domain
+      // metadata with empty clustering columns.
+      Seq.empty
+    }
+
+    filteredDomainMetadata ++ Seq(ClusteredTableUtils.createDomainMetadata(clusteringColumns))
   }
 
   /**

spark/src/main/scala/org/apache/spark/sql/delta/commands/RestoreTableCommand.scala

@@ -21,7 +21,7 @@ import java.sql.Timestamp
 import scala.collection.JavaConverters._
 import scala.util.{Success, Try}
 
-import org.apache.spark.sql.delta.{DeltaErrors, DeltaLog, DeltaOperations, Snapshot}
+import org.apache.spark.sql.delta.{DeltaErrors, DeltaLog, DeltaOperations, DomainMetadataUtils, Snapshot}
 import org.apache.spark.sql.delta.actions.{AddFile, DeletionVectorDescriptor, RemoveFile}
 import org.apache.spark.sql.delta.catalog.DeltaTableV2
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
@@ -205,9 +205,12 @@ case class RestoreTableCommand(sourceTable: DeltaTableV2)
           sourceProtocol.merge(targetProtocol)
         }
 
+        val actions = addActions ++ removeActions ++
+          DomainMetadataUtils.handleDomainMetadataForRestoreTable(snapshotToRestore, latestSnapshot)
+
         txn.commitLarge(
           spark,
-          addActions ++ removeActions,
+          actions,
           Some(newProtocol),
           DeltaOperations.Restore(version, timestamp),
           Map.empty,

spark/src/test/scala/org/apache/spark/sql/delta/skipping/ClusteredTableTestUtils.scala

@@ -152,7 +152,8 @@ trait ClusteredTableTestUtilsBase extends SparkFunSuite with SharedSparkSession
         } else {
           assertClusterByNotExist()
         }
-      case "WRITE" =>
+      case "WRITE" | "RESTORE" =>
+        // These are known operations from our tests that don't have clusterBy.
         doAssert(!lastOperationParameters.contains(CLUSTERING_PARAMETER_KEY))
       case _ =>
         // Other operations are not tested yet. If the test fails here, please check the expected

spark/src/test/scala/org/apache/spark/sql/delta/skipping/clustering/ClusteredTableDDLSuite.scala

@@ -21,6 +21,7 @@ import java.io.File
 import com.databricks.spark.util.{Log4jUsageLogger, MetricDefinitions}
 import org.apache.spark.sql.delta.skipping.ClusteredTableTestUtils
 import org.apache.spark.sql.delta.{DeltaAnalysisException, DeltaColumnMappingEnableIdMode, DeltaColumnMappingEnableNameMode, DeltaConfigs, DeltaExcludedBySparkVersionTestMixinShims, DeltaLog, DeltaUnsupportedOperationException}
+import org.apache.spark.sql.delta.clustering.ClusteringMetadataDomain
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
 import org.apache.spark.sql.delta.stats.SkippingEligibleDataType
 import org.apache.spark.sql.delta.test.{DeltaColumnMappingSelectedTestMixin, DeltaSQLCommandTest}
@@ -825,6 +826,45 @@ trait ClusteredTableDDLSuiteBase
     }
   }
 
+  test("validate RESTORE on clustered table") {
+    val tableIdentifier = TableIdentifier(testTable)
+    // Scenario 1: restore clustered table to unclustered version.
+    withTable(testTable) {
+      sql(s"CREATE TABLE $testTable (a INT, b STRING) USING delta")
+      val (_, startingSnapshot) = DeltaLog.forTableWithSnapshot(spark, tableIdentifier)
+      assert(!ClusteredTableUtils.isSupported(startingSnapshot.protocol))
+
+      sql(s"ALTER TABLE $testTable CLUSTER BY (a)")
+      verifyClusteringColumns(tableIdentifier, "a")
+
+      sql(s"RESTORE TABLE $testTable TO VERSION AS OF 0")
+      val (_, currentSnapshot) = DeltaLog.forTableWithSnapshot(spark, tableIdentifier)
+      verifyClusteringColumns(tableIdentifier, "")
+    }
+
+    // Scenario 2: restore clustered table to previous clustering columns.
+    withClusteredTable(testTable, "a INT, b STRING", "a") {
+      verifyClusteringColumns(tableIdentifier, "a")
+
+      sql(s"ALTER TABLE $testTable CLUSTER BY (b)")
+      verifyClusteringColumns(tableIdentifier, "b")
+
+      sql(s"RESTORE TABLE $testTable TO VERSION AS OF 0")
+      verifyClusteringColumns(tableIdentifier, "a")
+    }
+
+    // Scenario 3: restore unclustered table to clustered version.
+    withClusteredTable(testTable, "a int", "a") {
+      verifyClusteringColumns(tableIdentifier, "a")
+
+      sql(s"ALTER TABLE $testTable CLUSTER BY NONE")
+      verifyClusteringColumns(tableIdentifier, "")
+
+      sql(s"RESTORE TABLE $testTable TO VERSION AS OF 0")
+      verifyClusteringColumns(tableIdentifier, "a")
+    }
+  }
+
   testSparkMasterOnly("Variant is not supported") {
     val e = intercept[DeltaAnalysisException] {
       createOrReplaceClusteredTable("CREATE", testTable, "id long, v variant", "v")