[RLC] Case 4 — Capture & persist compaction OPTIMIZE source composition (write side)

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

@@ -1193,6 +1193,56 @@ case class RemoveFile(
 }
 // scalastyle:on
 
+object RemoveFile {
+  /**
+   * Misc tombstone-level metadata. Clients may safely ignore any of these tags; they must never
+   * affect correctness (an ignored tag only forgoes an optimization, e.g. a conflict reconcile).
+   */
+  object Tags {
+    /**
+     * [[COMPACTED_INTO]] / [[COMPACTION_INFO]]: recorded together on a source file removed by a
+     * compaction OPTIMIZE, describing where that source's rows landed in the compacted output, so
+     * the conflict checker can remap a concurrent deletion vector between the source and the output
+     * instead of aborting. The value format is modeled on what Databricks Runtime records;
+     * interoperating with a DBR-written OPTIMIZE on a shared table is best-effort and not a
+     * verified guarantee (an ignored or unrecognized tag only forgoes the reconcile -- never wrong
+     * data).
+     *
+     *  - [[COMPACTED_INTO]]: JSON array holding the single output path the source compacted into,
+     *    `["<output>.parquet"]` (matching the AddFile.path in the same commit).
+     *  - [[COMPACTION_INFO]]: JSON array holding the single run this source contributed,
+     *    `[{"rowOffsetInTarget": <outputStart>, "sourceNumPhysicalRecords": <physical>}]`. The
+     *    source's live rows land contiguously starting at physical offset `outputStart` of the
+     *    output, in source order. `sourceNumPhysicalRecords` is a PHYSICAL count; the live run
+     *    length is `sourceNumPhysicalRecords - |sourceDV|`, where `sourceDV` is the DV already on
+     *    this same tombstone (the DV the OPTIMIZE read). The physical count keeps the entry
+     *    self-consistent with the tombstone's own DV.
+     *
+     * Kept on the (short-lived) tombstone rather than the output AddFile (which snapshot
+     * reconstruction replays on every read); tombstone retention outlives the conflict window.
+     * Persisted (not stripped before commit) so a concurrent DML that LOSES to this OPTIMIZE can
+     * read the composition from the committed tombstone. Written whenever OPTIMIZE conflict
+     * reconciliation is enabled. O(1) per removed source.
+     */
+    val COMPACTED_INTO = "compactedInto"
+    val COMPACTION_INFO = "compactionInfo"
+  }
+}
+
+/**
+ * The per-source entry recorded in a compaction OPTIMIZE's `compactionInfo` tombstone tag: where
+ * that source's rows landed in the compacted output. The format is modeled on Databricks Runtime's
+ * (see [[RemoveFile.Tags.COMPACTION_INFO]]); reconciling against a DBR-written tag on a shared
+ * table is best-effort, not a verified guarantee. Every field is optional and unknown fields are
+ * ignored, so a foreign writer's schema drift degrades to a safe abort rather than a wrong result.
+ */
+@JsonIgnoreProperties(ignoreUnknown = true)
+private[delta] case class CompactionInfoEntry(
+    @JsonDeserialize(contentAs = classOf[java.lang.Long])
+    rowOffsetInTarget: Option[Long] = None,
+    @JsonDeserialize(contentAs = classOf[java.lang.Long])
+    sourceNumPhysicalRecords: Option[Long] = None)
+
 /**
  * A change file containing CDC data for the Delta version it's within. Non-CDC readers should
  * ignore this, CDC readers should scan all ChangeFiles in a version rather than computing

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

@@ -19,30 +19,33 @@ package org.apache.spark.sql.delta.commands
 import java.util.ConcurrentModificationException
 
 import scala.collection.mutable.ArrayBuffer
+import scala.util.control.NonFatal
 
 import org.apache.spark.sql.delta.skipping.MultiDimClustering
 import org.apache.spark.sql.delta.skipping.clustering.{ClusteredTableUtils, ClusteringColumnInfo}
 import org.apache.spark.sql.delta._
+import org.apache.spark.sql.delta.ClassicColumnConversions._
 import org.apache.spark.sql.delta.DeltaOperations.Operation
-import org.apache.spark.sql.delta.actions.{Action, AddFile, DeletionVectorDescriptor, FileAction, RemoveFile}
+import org.apache.spark.sql.delta.actions.{Action, AddFile, CompactionInfoEntry, DeletionVectorDescriptor, FileAction, RemoveFile}
 import org.apache.spark.sql.delta.commands.optimize._
-import org.apache.spark.sql.delta.files.SQLMetricsReporting
+import org.apache.spark.sql.delta.files.{SourceCompositionAccumulator, SourceCompositionCaptureExec, SQLMetricsReporting}
 import org.apache.spark.sql.delta.logging.DeltaLogKeys
 import org.apache.spark.sql.delta.schema.{SchemaUtils, UnsupportedDataTypeInfo}
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
-import org.apache.spark.sql.delta.util.BinPackingUtils
+import org.apache.spark.sql.delta.util.{BinPackingUtils, DeltaFileOperations, JsonUtils}
 
 import org.apache.spark.SparkContext
 import org.apache.spark.SparkContext.SPARK_JOB_GROUP_ID
 import org.apache.spark.internal.MDC
-import org.apache.spark.sql.{AnalysisException, Encoders, Row, SparkSession}
+import org.apache.spark.sql.{AnalysisException, DataFrame, Encoders, Row, SparkSession}
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedTable}
 import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Expression}
 import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnaryNode}
 import org.apache.spark.sql.execution.command.RunnableCommand
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.metric.SQLMetrics.createMetric
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 import org.apache.spark.util.{SystemClock, ThreadUtils}
 import org.apache.spark.sql.catalyst.catalog.CatalogTable
@@ -519,9 +522,40 @@ class OptimizeExecutor(
       bin: Seq[AddFile],
       maxFileSize: Long): Seq[FileAction] = {
     val baseTablePath = txn.deltaLog.dataPath
-
-    var input = txn.deltaLog.createDataFrame(txn.snapshot, bin, actionTypeOpt = Some("Optimize"))
+    // Compaction conflict-reconciliation (optimize.conflictReconciliation.enabled): observe the
+    // source composition of the compacted output at write time (SourceCompositionCaptureExec) so a
+    // concurrent DML's deletion vector can be remapped onto it instead of aborting. Compaction
+    // only, never clustering (a clustering pass permutes rows, so no offset mapping exists).
+    val reconcileEnabled = sparkSession.sessionState.conf
+      .getConf(DeltaSQLConf.DELTA_OPTIMIZE_CONFLICT_RECONCILIATION_ENABLED)
+    val useRepartition = sparkSession.sessionState.conf
+      .getConf(DeltaSQLConf.DELTA_OPTIMIZE_REPARTITION_ENABLED)
+    // Capture the source composition only on the coalesce path for compaction: coalesce
+    // preserves each source file's read order, so rows stay contiguous in the output
+    // (observed, not imposed -- no sort, no helper column) and row-range offsets exist.
+    // Repartition shuffles rows and a clustering pass permutes them, so neither is captured.
+    val captureReconcile =
+      reconcileEnabled && !isMultiDimClustering && !useRepartition
+
+    // Read the bin. The reconcile-capture path pins the read so each source file lands whole in
+    // one contiguous run (see readCompactionSourceWithWholeFilePins); vanilla OPTIMIZE just reads
+    // on the current session.
+    var input = if (captureReconcile) {
+      readCompactionSourceWithWholeFilePins(txn, bin, maxFileSize)
+    } else {
+      txn.deltaLog.createDataFrame(txn.snapshot, bin, actionTypeOpt = Some("Optimize"))
+    }
     input = RowTracking.preserveRowTrackingColumns(input, txn.snapshot)
+
+    val captureAccOpt: Option[SourceCompositionAccumulator] =
+      if (captureReconcile) {
+        val acc = new SourceCompositionAccumulator
+        sparkSession.sparkContext.register(acc)
+        Some(acc)
+      } else {
+        None
+      }
+
     val repartitionDF = if (isMultiDimClustering) {
       val totalSize = bin.map(_.size).sum
       val approxNumFiles = Math.max(1, totalSize / maxFileSize).toInt
@@ -531,8 +565,6 @@ class OptimizeExecutor(
         clusteringColumns,
         optimizeStrategy.curve)
     } else {
-      val useRepartition = sparkSession.sessionState.conf.getConf(
-        DeltaSQLConf.DELTA_OPTIMIZE_REPARTITION_ENABLED)
       if (useRepartition) {
         input.repartition(numPartitions = 1)
       } else {
@@ -549,18 +581,158 @@ class OptimizeExecutor(
       description)
 
     val binInfo = optimizeStrategy.initNewBin
-    val addFiles = txn.writeFiles(repartitionDF, None, isOptimize = true, Nil).collect {
+    val addFiles = txn.writeFiles(repartitionDF, None, isOptimize = true, Nil,
+        sourceCompositionCapture = captureAccOpt).collect {
       case a: AddFile => optimizeStrategy.tagAddFile(a, binInfo)
       case other =>
         throw new IllegalStateException(
           s"Unexpected action $other with type ${other.getClass}. File compaction job output" +
               s"should only have AddFiles")
     }
-    val removeFiles = bin.map(f => f.removeWithTimestamp(operationTimestamp, dataChange = false))
+    // Tombstones for the removed sources. Only the RLC reconciliation path tags each source with
+    // where its rows landed in the compacted output (so a concurrent DML's DV can be remapped by
+    // offset instead of aborting); vanilla OPTIMIZE uses plain untagged tombstones and skips the
+    // reconcile helper entirely.
+    val removeFiles = if (captureReconcile) {
+      buildRemoveFilesWithCompactionCompositionTags(
+        txn, bin, addFiles, captureAccOpt, operationTimestamp)
+    } else {
+      bin.map(_.removeWithTimestamp(operationTimestamp, dataChange = false))
+    }
     val updates = addFiles ++ removeFiles
     updates
   }
 
+  /**
+   * Read `bin` for a compaction OPTIMIZE on the reconciliation-capture path, pinning the read so
+   * each source file lands whole in a single partition.
+   *
+   * Each source file must be read whole so no source is split across partitions -- coalesce(1)
+   * then lands each source as one contiguous run, which is what the capture's one-run-per-file
+   * gate needs. Spark has no "do not split" toggle for Parquet; splitting is governed by
+   *   maxSplitBytes = min(maxPartitionBytes, max(openCostInBytes, totalBytes / minPartitionNum))
+   * so pin BOTH maxPartitionBytes (>= the compaction target) and minPartitionNum = 1: then
+   * maxSplitBytes >= every source (each <= the target), so no source splits. (A split that still
+   * somehow slips through just fails the capture gate and aborts -- never wrong data.) Both confs
+   * go on a CLONED session so the override is isolated from other queries -- and from the other
+   * bins compacting concurrently -- that share this SparkSession: createDataFrame binds the scan
+   * relation to SparkSession.active and Spark reads these confs from the captured session, so the
+   * clone need only be active while the relation is built, then the previous active session is
+   * restored.
+   */
+  private def readCompactionSourceWithWholeFilePins(
+      txn: OptimisticTransaction,
+      bin: Seq[AddFile],
+      maxFileSize: Long): DataFrame = {
+    val readSession = sparkSession.cloneSession()
+    readSession.conf.set(SQLConf.FILES_MAX_PARTITION_BYTES.key, maxFileSize)
+    readSession.conf.set(SQLConf.FILES_MIN_PARTITION_NUM.key, "1")
+    val prevActive = SparkSession.getActiveSession
+    SparkSession.setActiveSession(readSession)
+    try {
+      txn.deltaLog.createDataFrame(txn.snapshot, bin, actionTypeOpt = Some("Optimize"))
+    } finally {
+      prevActive.fold(SparkSession.clearActiveSession())(SparkSession.setActiveSession)
+    }
+  }
+
+  /**
+   * Build the removed-source tombstones for a compaction OPTIMIZE, tagging each with its
+   * `compactedInto` / `compactionInfo` composition -- where that source's rows landed in the single
+   * compacted output -- so a concurrent DML's deletion vector can be remapped by offset instead of
+   * aborting (see [[RemoveFile.Tags.COMPACTION_INFO]]). The composition is persisted, as Databricks
+   * Runtime does on every compaction OPTIMIZE: consumed in-memory when THIS OPTIMIZE loses to a
+   * concurrent DML, and read back from the committed tombstone when a concurrent DML LOSES to this
+   * OPTIMIZE. The tag format is modeled on the one Databricks Runtime writes; reconciling against a
+   * DBR-written tag on a shared table is best-effort, not a verified guarantee.
+   *
+   * Falls back to plain untagged tombstones -- so the conflict aborts as it does today -- unless
+   * the capture is present (reconciliation was enabled) AND trustworthy: exactly one output file
+   * from exactly one write partition, and each source contributed exactly one captured run covering
+   * the whole bin (a sanity gate against retries / speculation / splits).
+   *
+   * Each source's `compactionInfo` records `sourceNumPhysicalRecords` (the live rows the write saw
+   * PLUS the source's read-time DV cardinality), not the live count, so the tags stay O(1) per
+   * source regardless of how fragmented a source DV is; the conflict checker recovers the live run
+   * length by subtracting the tombstone's own DV, and rebuilds the read-time gaps only on a real
+   * conflict.
+   */
+  private def buildRemoveFilesWithCompactionCompositionTags(
+      txn: OptimisticTransaction,
+      bin: Seq[AddFile],
+      addFiles: Seq[AddFile],
+      captureAccOpt: Option[SourceCompositionAccumulator],
+      operationTimestamp: Long): Seq[RemoveFile] = {
+    // The fallback: plain untagged tombstones, exactly as vanilla OPTIMIZE writes them, whenever
+    // the capture is off or the trustworthiness gate below rejects it (the loser aborts as today).
+    def untagged: Seq[RemoveFile] =
+      bin.map(_.removeWithTimestamp(operationTimestamp, dataChange = false))
+
+    try captureAccOpt match {
+      case Some(acc) if addFiles.size == 1 && acc.value.size() == 1 &&
+          bin.forall(_.numLogicalRecords.isDefined) =>
+        val runs = acc.value.get(0)
+        val captured = (0 until runs.size()).map(runs.get(_).count).sum
+        val expected = bin.flatMap(_.numLogicalRecords).sum
+        // Map each run's absolute source path (from the holder) back to the table-relative AddFile
+        // path, so the recorded keys match the RemoveFiles / winning DV updates (both relative) at
+        // conflict-resolution time.
+        val nameToAddFile = generateCandidateFileMap(txn.deltaLog.dataPath, bin)
+        val tablePath = txn.deltaLog.dataPath
+        val compactedIntoJson = JsonUtils.toJson(Seq(addFiles.head.path))
+        // Walk runs in output (write) order, accumulating each source's start offset in the output.
+        var outputPos = 0L
+        val tagByPath: Seq[Option[(String, Map[String, String])]] = (0 until runs.size()).map { i =>
+          val r = runs.get(i)
+          val start = outputPos
+          outputPos += r.count
+          // A null source file (the holder was empty for some rows) can't be mapped back to an
+          // AddFile; yield None so the whole capture is treated as unreconcilable below (no NPE in
+          // absolutePath), and the loser aborts as it does today.
+          if (r.sourceFile == null) {
+            None
+          } else {
+            val abs =
+              DeltaFileOperations.absolutePath(tablePath.toString, r.sourceFile).toString
+            nameToAddFile.get(abs).map { add =>
+              // `r.count` is the live rows the write saw; the physical count adds back the source's
+              // read-time DV (the DV carried onto this source's tombstone below).
+              val sourceDvCardinality = Option(add.deletionVector).map(_.cardinality).getOrElse(0L)
+              val physical = r.count + sourceDvCardinality
+              val compactionInfoJson =
+                JsonUtils.toJson(Seq(CompactionInfoEntry(Some(start), Some(physical))))
+              add.path -> Map(
+                RemoveFile.Tags.COMPACTED_INTO -> compactedIntoJson,
+                RemoveFile.Tags.COMPACTION_INFO -> compactionInfoJson)
+            }
+          }
+        }
+        // Each source file mapped and produced exactly one captured run covering the whole bin.
+        val oneRunPerFile = runs.size() == bin.size &&
+          (0 until runs.size()).map(runs.get(_).sourceFile).distinct.size == bin.size
+        if (captured == expected && tagByPath.forall(_.isDefined) && oneRunPerFile) {
+          val tags = tagByPath.flatten.toMap
+          bin.map { f =>
+            val r = f.removeWithTimestamp(operationTimestamp, dataChange = false)
+            tags.get(f.path).fold(r)(
+              _.foldLeft(r) { case (tagged, (k, v)) => tagged.copyWithTag(k, v) })
+          }
+        } else {
+          untagged
+        }
+      case _ =>
+        untagged
+    } catch {
+      case NonFatal(e) =>
+        // Composition capture is a pure optimization enabler, never required for OPTIMIZE
+        // correctness: on any failure building the tags (an unmappable source path, a
+        // serialization error) fall back to plain untagged tombstones so the already-written
+        // OPTIMIZE still commits and a concurrent loser aborts exactly as it does today.
+        logWarning(log"Compaction composition capture failed; writing untagged tombstones", e)
+        untagged
+    }
+  }
+
   /**
    * Attempts to commit the given actions to the log. In the case of a concurrent update,
    * the given function will be invoked with a new transaction to allow custom conflict

spark/src/main/scala/org/apache/spark/sql/delta/files/DeltaFileFormatWriter.scala

@@ -61,12 +61,12 @@ object DeltaFileFormatWriter extends Logging {
    * A variable used in tests to check whether the output ordering of the query matches the
    * required ordering of the write command.
    */
-  private var outputOrderingMatched: Boolean = false
+  private[delta] var outputOrderingMatched: Boolean = false
 
   /**
    * A variable used in tests to check the final executed plan.
    */
-  private var executedPlan: Option[SparkPlan] = None
+  private[delta] var executedPlan: Option[SparkPlan] = None
 
   // scalastyle:off argcount
   /**