[Spark] REPLACE PARTITIONED BY: rewrite data files that do not store partition columns

spark/src/main/resources/error/delta-error-classes.json

@@ -2763,7 +2763,7 @@
   "DELTA_REPLACE_PARTITIONED_BY_PARTITION_COLUMNS_NOT_MATERIALIZED" : {
     "message" : [
       "Cannot replace the partitioning of the table because the partition column(s) <partitionColumns> are not stored in <numFiles> data file(s), for example: <examplePaths>.",
-      "Rewrite these files so that they contain the partition column values before running ALTER TABLE REPLACE PARTITIONED BY."
+      "Enable <config> to rewrite these files as part of the command, or rewrite them so that they contain the partition column values before running ALTER TABLE REPLACE PARTITIONED BY."
     ],
     "sqlState" : "55000"
   },

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

@@ -4284,7 +4284,10 @@ trait DeltaErrorsBase
     new DeltaAnalysisException(
       errorClass = "DELTA_REPLACE_PARTITIONED_BY_PARTITION_COLUMNS_NOT_MATERIALIZED",
       messageParameters = Array(
-        partitionColumns.mkString(", "), numFiles.toString, examplePaths.mkString(", ")))
+        partitionColumns.mkString(", "),
+        numFiles.toString,
+        examplePaths.mkString(", "),
+        DeltaSQLConf.DELTA_REPLACE_PARTITIONED_BY_REWRITE_NON_MATERIALIZED_FILES.key))
   }
 
   def createTableWithDifferentClusteringException(

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

@@ -24,7 +24,7 @@ import scala.util.Try
 import org.apache.spark.sql.delta.skipping.clustering.ClusteredTableUtils
 import org.apache.spark.sql.delta.skipping.clustering.temp.ClusterBySpec
 import org.apache.spark.sql.delta._
-import org.apache.spark.sql.delta.actions.Action
+import org.apache.spark.sql.delta.actions.{Action, AddFile, FileAction}
 import org.apache.spark.sql.delta.catalog.DeltaTableV2
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
 import org.apache.spark.sql.delta.stats.DeltaStatistics
@@ -41,23 +41,27 @@ import org.apache.spark.sql.types._
 import org.apache.spark.util.SerializableConfiguration
 
 /**
- * Command that replaces the partitioning of a Delta table with clustering, without rewriting
- * any data file:
+ * Command that replaces the partitioning of a Delta table with clustering:
  *  - ALTER TABLE .. REPLACE PARTITIONED BY WITH CLUSTER BY (col1, col2, ...)
  *  - ALTER TABLE .. REPLACE PARTITIONED BY WITH CLUSTER BY NONE
  *
- * The conversion is a single metadata-only commit which
- *  1. clears the partition columns in the table metadata,
- *  2. re-adds every active file with empty `partitionValues` and `dataChange = false`, keeping its
- *     path, deletion vector and row tracking fields. Per-file min/max/nullCount statistics are
- *     synthesized for the former partition columns from the (constant) partition values so that
- *     data skipping on these columns keeps working, and
- *  3. when clustering columns are given, enables clustering and records the clustering columns.
- *
  * After the conversion the former partition columns are read from the data files, so they must be
  * physically stored in every data file (see `delta.writePartitionColumnsToParquet` and the
- * `materializePartitionColumns` table feature). This is verified by reading the Parquet footers
- * before committing.
+ * `materializePartitionColumns` table feature). The command reads the Parquet footers of all data
+ * files to find the files that do not store them.
+ *
+ * The conversion is a single commit which
+ *  1. clears the partition columns in the table metadata,
+ *  2. re-adds every file that stores the partition columns with empty `partitionValues` and
+ *     `dataChange = false`, keeping its path, deletion vector and row tracking fields. Per-file
+ *     min/max/nullCount statistics are synthesized for the former partition columns from the
+ *     (constant) partition values so that data skipping on these columns keeps working,
+ *  3. rewrites every file that does not store the partition columns (if any), like OPTIMIZE does:
+ *     the new files store all columns, deletion vectors are applied and row IDs and row commit
+ *     versions are preserved, and
+ *  4. when clustering columns are given, enables clustering and records the clustering columns.
+ *
+ * If all files store the partition columns, the conversion does not rewrite any data.
  *
  * The commit is done with [[OptimisticTransaction.commitLarge]], which fails if any other commit
  * happened since the transaction started. Hence no file can be added concurrently between the
@@ -80,11 +84,20 @@ case class AlterTableReplacePartitionedByWithClusterByDeltaCommand(
       ClusteredTableUtils.validateNumClusteringColumns(clusteringColumns, Some(deltaLog))
 
       val physicalPartitionSchema = oldMetadata.physicalPartitionSchema
-      if (sparkSession.sessionState.conf.getConf(
+      val conf = sparkSession.sessionState.conf
+      val filesToRewrite = if (conf.getConf(
           DeltaSQLConf.DELTA_REPLACE_PARTITIONED_BY_VERIFY_MATERIALIZED_PARTITION_COLUMNS)) {
-        ReplacePartitionedByUtils.verifyPartitionColumnsMaterialized(
-          sparkSession, deltaLog, snapshot, oldMetadata.partitionColumns,
-          physicalPartitionSchema.fieldNames.toSeq)
+        ReplacePartitionedByUtils.findFilesWithoutPartitionColumns(
+          sparkSession, deltaLog, snapshot, physicalPartitionSchema.fieldNames.toSeq)
+      } else {
+        Seq.empty
+      }
+      if (filesToRewrite.nonEmpty && !conf.getConf(
+          DeltaSQLConf.DELTA_REPLACE_PARTITIONED_BY_REWRITE_NON_MATERIALIZED_FILES)) {
+        throw DeltaErrors.replacePartitionedByPartitionColumnsNotMaterializedException(
+          oldMetadata.partitionColumns,
+          filesToRewrite.size,
+          filesToRewrite.take(ReplacePartitionedByUtils.NUM_EXAMPLE_FILES).map(_.path))
       }
 
       val newLogicalClusteringColumns = clusteringColumns.map(FieldReference(_).toString)
@@ -106,12 +119,17 @@ case class AlterTableReplacePartitionedByWithClusterByDeltaCommand(
         Nil
       }
 
-      val stringPrefixLength = sparkSession.sessionState.conf.getConf(
-        DeltaSQLConf.DATA_SKIPPING_STRING_PREFIX_LENGTH)
+      // Must run after the metadata update, so that the new files store all columns.
+      val rewriteActions =
+        ReplacePartitionedByUtils.rewriteFiles(sparkSession, txn, snapshot, filesToRewrite)
+
+      val stringPrefixLength = conf.getConf(DeltaSQLConf.DATA_SKIPPING_STRING_PREFIX_LENGTH)
       val partitionFields = physicalPartitionSchema.fields.toSeq
+      val rewrittenPaths = sparkSession.sparkContext.broadcast(filesToRewrite.map(_.path).toSet)
       import org.apache.spark.sql.delta.implicits._
       val newAddFiles = snapshot.allFiles.mapPartitions { files =>
-        files.map { file =>
+        val skip = rewrittenPaths.value
+        files.filterNot(file => skip.contains(file.path)).map { file =>
           file.copy(
             partitionValues = Map.empty,
             dataChange = false,
@@ -121,7 +139,12 @@ case class AlterTableReplacePartitionedByWithClusterByDeltaCommand(
       }
 
       val actions: Iterator[Action] =
-        newAddFiles.toLocalIterator().asScala ++ domainMetadata.iterator
+        rewriteActions.iterator ++ newAddFiles.toLocalIterator().asScala ++ domainMetadata.iterator
+      val metrics = Map(
+        "numRewrittenFiles" -> filesToRewrite.size.toString,
+        "numRewrittenBytes" -> filesToRewrite.map(_.size).sum.toString,
+        "numAddedFilesFromRewrite" ->
+          rewriteActions.count(_.isInstanceOf[AddFile]).toString)
       val newProtocolOpt =
         if (txn.protocol != snapshot.protocol) Some(txn.protocol) else None
       txn.commitLarge(
@@ -132,7 +155,7 @@ case class AlterTableReplacePartitionedByWithClusterByDeltaCommand(
           oldMetadata.partitionColumns.mkString(","),
           newLogicalClusteringColumns.mkString(",")),
         context = Map.empty,
-        metrics = Map.empty,
+        metrics = metrics,
         dataChange = Some(false))
     }
     Seq.empty[Row]
@@ -141,42 +164,62 @@ case class AlterTableReplacePartitionedByWithClusterByDeltaCommand(
 
 object ReplacePartitionedByUtils {
 
-  private val NUM_EXAMPLE_FILES = 5
+  val NUM_EXAMPLE_FILES = 5
 
   /**
-   * Verifies that every active data file of the snapshot physically stores all the given
+   * Returns the active data files of the snapshot that do not physically store all the given
    * partition columns, by reading the Parquet footers of the files.
    */
-  def verifyPartitionColumnsMaterialized(
+  def findFilesWithoutPartitionColumns(
       spark: SparkSession,
       deltaLog: DeltaLog,
       snapshot: Snapshot,
-      logicalPartitionColumns: Seq[String],
-      physicalPartitionColumns: Seq[String]): Unit = {
+      physicalPartitionColumns: Seq[String]): Seq[AddFile] = {
     val broadcastConf = spark.sparkContext.broadcast(
       new SerializableConfiguration(deltaLog.newDeltaHadoopConf()))
     val dataRootDir = deltaLog.dataPath.toString
     val requiredColumns = physicalPartitionColumns.map(_.toLowerCase(Locale.ROOT))
 
     import org.apache.spark.sql.delta.implicits._
-    val filesMissingColumns = snapshot.allFiles.select("path").as[String].mapPartitions { paths =>
+    snapshot.allFiles.mapPartitions { files =>
       val conf = broadcastConf.value.value
-      paths.filter { relativePath =>
-        val path = DeltaFileOperations.absolutePath(dataRootDir, relativePath)
+      files.filter { file =>
+        val path = DeltaFileOperations.absolutePath(dataRootDir, file.path)
         val status = path.getFileSystem(conf).getFileStatus(path)
         val footer = ParquetFooterReaderShims.readParquetFooter(
           conf, status, ParquetMetadataConverter.SKIP_ROW_GROUPS)
         val fileColumns = footer.getFileMetaData.getSchema.getFields.asScala
           .map(_.getName.toLowerCase(Locale.ROOT)).toSet
         !requiredColumns.forall(fileColumns.contains)
       }
-    }
+    }.collect().toSeq
+  }
 
-    val examples = filesMissingColumns.take(NUM_EXAMPLE_FILES)
-    if (examples.nonEmpty) {
-      throw DeltaErrors.replacePartitionedByPartitionColumnsNotMaterializedException(
-        logicalPartitionColumns, filesMissingColumns.count(), examples.toSeq)
+  /**
+   * Rewrites the given files with the metadata of the transaction, which must not have partition
+   * columns anymore, so that the new files store all columns. Deletion vectors are applied, and
+   * row IDs and row commit versions are preserved like in OPTIMIZE. Returns the new files and the
+   * removal of the given files, all with `dataChange = false`.
+   */
+  def rewriteFiles(
+      spark: SparkSession,
+      txn: OptimisticTransaction,
+      snapshot: Snapshot,
+      files: Seq[AddFile]): Seq[FileAction] = {
+    if (files.isEmpty) return Seq.empty
+    require(txn.metadata.partitionColumns.isEmpty,
+      "The partitioning must be dropped before rewriting the files")
+    var input = txn.deltaLog.createDataFrame(
+      snapshot, files, actionTypeOpt = Some("ReplacePartitionedBy"))
+    input = RowTracking.preserveRowTrackingColumns(input, snapshot)
+    val addFiles = txn.writeFiles(input, None, isOptimize = true, Nil).collect {
+      case a: AddFile => a.copy(dataChange = false)
+      case other =>
+        throw new IllegalStateException(
+          s"Unexpected action $other with type ${other.getClass} while rewriting files")
     }
+    val timestamp = System.currentTimeMillis()
+    addFiles ++ files.map(_.removeWithTimestamp(timestamp, dataChange = false))
   }
 
   /**

spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSQLConf.scala

@@ -3455,6 +3455,16 @@ trait DeltaSQLConfBase extends DeltaSQLConfUtils {
       .booleanConf
       .createWithDefault(true)
 
+  val DELTA_REPLACE_PARTITIONED_BY_REWRITE_NON_MATERIALIZED_FILES =
+    buildConf("alterTable.replacePartitionedBy.rewriteNonMaterializedFiles")
+      .internal()
+      .doc("""When true, ALTER TABLE ... REPLACE PARTITIONED BY WITH CLUSTER BY rewrites the data
+        |files that do not physically store the partition columns, so that the partition column
+        |values are stored in the new files. When false, the command fails if such files exist.
+        |""".stripMargin)
+      .booleanConf
+      .createWithDefault(true)
+
   val DELTA_LOG_CACHE_SIZE = buildConf("delta.log.cacheSize")
     .internal()
     .doc("The maximum number of DeltaLog instances to cache in memory.")

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

@@ -19,7 +19,7 @@ package org.apache.spark.sql.delta.skipping.clustering
 import org.apache.spark.sql.delta.skipping.ClusteredTableTestUtils
 import org.apache.spark.sql.delta.skipping.clustering.temp.{AlterTableReplacePartitionedByWithClusterBy, ClusterBySpec}
 import org.apache.spark.sql.delta.{DeltaAnalysisException, DeltaLog}
-import org.apache.spark.sql.delta.actions.AddFile
+import org.apache.spark.sql.delta.actions.{AddFile, RemoveFile}
 import org.apache.spark.sql.delta.commands.ReplacePartitionedByUtils
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
 import org.apache.spark.sql.delta.stats.DataSkippingDeltaTestsUtils
@@ -92,6 +92,8 @@ class ReplacePartitionedByWithClusterBySuite
 
       val changes = deltaLog.getChanges(after.version).next()._2
       assert(changes.collect { case a: AddFile => a }.forall(!_.dataChange))
+      assert(!changes.exists(_.isInstanceOf[RemoveFile]))
+      assert(lastCommit.operationMetrics.flatMap(_.get("numRewrittenFiles")) === Some("0"))
 
       // Time travel to the partitioned version still works.
       checkAnswer(
@@ -245,20 +247,119 @@ class ReplacePartitionedByWithClusterBySuite
     }
   }
 
-  test("fails when data files do not store the partition columns") {
+  private def filesWithoutPartitionColumns(path: String, partitionColumns: String*) = {
+    val deltaLog = DeltaLog.forTable(spark, path)
+    ReplacePartitionedByUtils.findFilesWithoutPartitionColumns(
+      spark, deltaLog, deltaLog.update(), partitionColumns)
+  }
+
+  private def lastOperationMetrics(path: String): Map[String, String] =
+    sql(s"DESCRIBE HISTORY delta.`$path` LIMIT 1")
+      .select("operationMetrics").head().getMap[String, String](0).toMap
+
+  test("rewrites the data files that do not store the partition columns") {
+    withTempDir { dir =>
+      val path = dir.getCanonicalPath
+      sql(s"""CREATE TABLE delta.`$path` (id BIGINT, p STRING) USING delta PARTITIONED BY (p)
+             |TBLPROPERTIES ('delta.writePartitionColumnsToParquet' = 'false')""".stripMargin)
+      sql(s"INSERT INTO delta.`$path` VALUES (1, 'a'), (2, 'b')")
+      sql(s"ALTER TABLE delta.`$path` SET TBLPROPERTIES " +
+        "('delta.writePartitionColumnsToParquet' = 'true')")
+      sql(s"INSERT INTO delta.`$path` VALUES (3, 'a'), (4, null)")
+      val deltaLog = DeltaLog.forTable(spark, path)
+      val before = deltaLog.update()
+      val notMaterialized = filesWithoutPartitionColumns(path, "p").map(_.path).toSet
+      assert(notMaterialized.size === 2)
+      val materialized = before.allFiles.collect().map(_.path).toSet -- notMaterialized
+      assert(materialized.size === 2)
+
+      sql(s"ALTER TABLE delta.`$path` REPLACE PARTITIONED BY WITH CLUSTER BY NONE")
+
+      val after = deltaLog.update()
+      assert(after.version === before.version + 1)
+      assert(after.metadata.partitionColumns.isEmpty)
+      val afterFiles = after.allFiles.collect()
+      assert(afterFiles.forall(_.partitionValues.isEmpty))
+      val afterPaths = afterFiles.map(_.path).toSet
+      assert(materialized.subsetOf(afterPaths))
+      assert((afterPaths intersect notMaterialized).isEmpty)
+      assert(filesWithoutPartitionColumns(path, "p").isEmpty)
+      // Synthesized stats for kept files, collected stats for rewritten files.
+      afterFiles.foreach { f =>
+        assert(statsOf(f)("nullCount").asInstanceOf[Map[String, Any]].contains("p"))
+      }
+      val rewrittenStats = afterFiles.filterNot(f => materialized.contains(f.path)).map(statsOf)
+      assert(rewrittenStats.nonEmpty)
+      def pStat(kind: String): Seq[String] =
+        rewrittenStats.map(_(kind).asInstanceOf[Map[String, Any]]("p").toString).toSeq
+      assert(pStat("minValues").min === "a")
+      assert(pStat("maxValues").max === "b")
+      assert(filesRead(spark, deltaLog, "p = 'b'", checkEmptyUnusedFilters = false) <
+        afterFiles.length)
+
+      val changes = deltaLog.getChanges(after.version).next()._2
+      val removed = changes.collect { case r: RemoveFile => r }
+      assert(removed.map(_.path).toSet === notMaterialized)
+      assert(removed.forall(!_.dataChange))
+      assert(changes.collect { case a: AddFile => a }.forall(!_.dataChange))
+      val metrics = lastOperationMetrics(path)
+      assert(metrics("numRewrittenFiles") === "2")
+      assert(metrics("numAddedFilesFromRewrite").toInt > 0)
+
+      checkAnswer(readTable(path),
+        Seq(Row(1L, "a"), Row(2L, "b"), Row(3L, "a"), Row(4L, null)))
+      checkAnswer(readTable(path).where("p = 'b'").select("id"), Row(2L))
+      checkAnswer(readTable(path).where("p IS NULL").select("id"), Row(4L))
+    }
+  }
+
+  test("rewrite applies deletion vectors and preserves row tracking") {
+    withTempDir { dir =>
+      val path = dir.getCanonicalPath
+      sql(s"""CREATE TABLE delta.`$path` (id BIGINT, p STRING) USING delta PARTITIONED BY (p)
+             |TBLPROPERTIES (
+             |  'delta.writePartitionColumnsToParquet' = 'false',
+             |  'delta.enableDeletionVectors' = 'true',
+             |  'delta.enableRowTracking' = 'true')""".stripMargin)
+      spark.range(0, 100, 1, 4).selectExpr("id", "cast(id % 3 as string) as p")
+        .write.format("delta").mode("append").save(path)
+      sql(s"DELETE FROM delta.`$path` WHERE id % 10 = 0")
+      val deltaLog = DeltaLog.forTable(spark, path)
+      assert(deltaLog.update().allFiles.collect().exists(_.deletionVector != null))
+      val rowTrackingColumns =
+        Seq("*", "_metadata.row_id", "_metadata.row_commit_version")
+      val expected = readTable(path).selectExpr(rowTrackingColumns: _*).collect()
+      assert(expected.length === 90)
+
+      sql(s"ALTER TABLE delta.`$path` REPLACE PARTITIONED BY WITH CLUSTER BY NONE")
+
+      val after = deltaLog.update()
+      assert(after.metadata.partitionColumns.isEmpty)
+      assert(after.allFiles.collect().forall(_.deletionVector == null))
+      assert(filesWithoutPartitionColumns(path, "p").isEmpty)
+      checkAnswer(readTable(path).selectExpr(rowTrackingColumns: _*), expected)
+    }
+  }
+
+  test("fails when data files do not store the partition columns and rewrite is disabled") {
     withTempDir { dir =>
       val path = dir.getCanonicalPath
       sql(s"""CREATE TABLE delta.`$path` (id BIGINT, p STRING) USING delta PARTITIONED BY (p)
              |TBLPROPERTIES ('delta.writePartitionColumnsToParquet' = 'false')""".stripMargin)
       sql(s"INSERT INTO delta.`$path` VALUES (1, 'a'), (2, 'b')")
       val deltaLog = DeltaLog.forTable(spark, path)
       val versionBefore = deltaLog.update().version
-      val e = intercept[DeltaAnalysisException] {
-        sql(s"ALTER TABLE delta.`$path` REPLACE PARTITIONED BY WITH CLUSTER BY NONE")
+      val e = withSQLConf(DeltaSQLConf
+          .DELTA_REPLACE_PARTITIONED_BY_REWRITE_NON_MATERIALIZED_FILES.key -> "false") {
+        intercept[DeltaAnalysisException] {
+          sql(s"ALTER TABLE delta.`$path` REPLACE PARTITIONED BY WITH CLUSTER BY NONE")
+        }
       }
       assert(e.getErrorClass ===
         "DELTA_REPLACE_PARTITIONED_BY_PARTITION_COLUMNS_NOT_MATERIALIZED")
       assert(e.getMessage.contains("2 data file(s)"))
+      assert(e.getMessage.contains(
+        DeltaSQLConf.DELTA_REPLACE_PARTITIONED_BY_REWRITE_NON_MATERIALIZED_FILES.key))
       assert(deltaLog.update().version === versionBefore)
       assert(deltaLog.update().metadata.partitionColumns === Seq("p"))
     }