[Spark] Add invariant checks to DML commands (#4486)

<!--
Thanks for sending a pull request!  Here are some tips for you:
1. If this is your first time, please read our contributor guidelines:
https://github.com/delta-io/delta/blob/master/CONTRIBUTING.md
2. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP]
Your PR title ...'.
  3. Be sure to keep the PR description updated to reflect all changes.
  4. Please write your PR title to summarize what this PR proposes.
5. If possible, provide a concise example to reproduce the issue for a
faster review.
6. If applicable, include the corresponding issue number in the PR title
and link it in the body.
-->

#### Which Delta project/connector is this regarding?
<!--
Please add the component selected below to the beginning of the pull
request title
For example: [Spark] Title of my pull request
-->

- [x] Spark
- [ ] Standalone
- [ ] Flink
- [ ] Kernel
- [ ] Other (fill in here)

## Description

This PR adds invariant checks to DELETE, UPDATE, and MERGE to detect
bugs in Spark and Delta, and to prevent these DML statements from
committing if they suffered from the bug.

These invariants come in two flavors:
- Unreliable checks using the commit stats (i.e. number of rows in the
files added and removed) and the SQL metrics. These checks are disabled
by default, as Spark can overcount metrics when there are retries.
- Reliable checks based purely on the commit stats. These checks are
enabled by default, but cannot detect every occurrence of a bug.

## How was this patch tested?

- Manually ran existing DML tests with all invariant checks enabled by
default to confirm that the checks do not cause any issues.
- Manually ran existing DML tests with all invariant checks enabled by
default with a bug introduced to confirm that the checks trigger.
- Added `DeltaCommandInvariantsSuite`.

## Does this PR introduce _any_ user-facing changes?

<!--
If yes, please clarify the previous behavior and the change this PR
proposes - provide the console output, description and/or an example to
show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change
compared to the released Delta Lake versions or within the unreleased
branches such as master.
If no, write 'No'.
-->
No

Author: tomvanbussel

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

@@ -546,6 +546,14 @@
     ],
     "sqlState" : "2200G"
   },
+  "DELTA_COMMAND_INVARIANT_VIOLATION" : {
+    "message" : [
+      "A command internal invariant was violated in '<operation>'.",
+      "Please retry the command.",
+      "Exception reference: <uuid>."
+    ],
+    "sqlState" : "XXKDS"
+  },
   "DELTA_COMMIT_INTERMEDIATE_REDIRECT_STATE" : {
     "message" : [
       "Cannot handle commit of table within redirect table state '<state>'."
@@ -1994,6 +2002,14 @@
     ],
     "sqlState" : "42P18"
   },
+  "DELTA_NUM_RECORDS_MISMATCH" : {
+    "message" : [
+      "Failed to validate the number of records in <operation>.",
+      "Added <numAddedRecords> records and removed <numRemovedRecords> records.",
+      "This is a bug."
+    ],
+    "sqlState" : "XXKDS"
+  },
   "DELTA_ONEOF_IN_TIMETRAVEL" : {
     "message" : [
       "Please either provide 'timestampAsOf' or 'versionAsOf' for time travel."

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

@@ -19,7 +19,7 @@ package org.apache.spark.sql.delta
 // scalastyle:off import.ordering.noEmptyLine
 import java.io.{FileNotFoundException, IOException}
 import java.nio.file.FileAlreadyExistsException
-import java.util.ConcurrentModificationException
+import java.util.{ConcurrentModificationException, UUID}
 
 import scala.collection.JavaConverters._
 
@@ -3728,6 +3728,25 @@ trait DeltaErrorsBase
       errorClass = "DELTA_UNSUPPORTED_CATALOG_OWNED_TABLE_CREATION",
       messageParameters = Array.empty)
   }
+
+  def numRecordsMismatch(
+      operation: String,
+      numAddedRecords: Long,
+      numRemovedRecords: Long): Throwable = {
+    new DeltaIllegalStateException(
+      errorClass = "DELTA_NUM_RECORDS_MISMATCH",
+      messageParameters = Array(operation, numAddedRecords.toString, numRemovedRecords.toString)
+    )
+  }
+
+  def commandInvariantViolationException(
+      operation: String,
+      id: UUID): Throwable = {
+    new DeltaIllegalStateException(
+      errorClass = "DELTA_COMMAND_INVARIANT_VIOLATION",
+      messageParameters = Array(operation, id.toString)
+    )
+  }
 }
 
 object DeltaErrors extends DeltaErrorsBase

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

@@ -30,7 +30,8 @@ case class NumRecordsStats (
     numDeletionVectorRecordsAdded: Long,
     numDeletionVectorRecordsRemoved: Long,
     numFilesAddedWithoutNumRecords: Long,
-    numFilesRemovedWithoutNumRecords: Long) {
+    numFilesRemovedWithoutNumRecords: Long,
+    numLogicalRecordsAddedInFilesWithDeletionVectorsPartial: Long) {
 
   def allFilesHaveNumRecords: Boolean =
     numFilesAddedWithoutNumRecords == 0 && numFilesRemovedWithoutNumRecords == 0
@@ -48,6 +49,14 @@ case class NumRecordsStats (
    */
   def numLogicalRecordsRemoved: Option[Long] = Option.when(numFilesRemovedWithoutNumRecords == 0)(
     numLogicalRecordsRemovedPartial)
+
+  /**
+   * The number of logical records in all AddFile actions that have a deletion vector or None
+   * if any file does not contain statistics.
+   */
+  def numLogicalRecordsAddedInFilesWithDeletionVectors: Option[Long] =
+    Option.when(numFilesAddedWithoutNumRecords == 0)(
+      numLogicalRecordsAddedInFilesWithDeletionVectorsPartial)
 }
 
 object NumRecordsStats {
@@ -60,6 +69,7 @@ object NumRecordsStats {
     var numLogicalRecordsRemovedPartial: Long = 0L
     var numDeletionVectorRecordsAdded = 0L
     var numDeletionVectorRecordsRemoved = 0L
+    var numLogicalRecordsAddedInFilesWithDeletionVectorsPartial = 0L
 
     actions.foreach {
       case a: AddFile =>
@@ -69,6 +79,10 @@ object NumRecordsStats {
           0L
         }
         numDeletionVectorRecordsAdded += a.numDeletedRecords
+        if (a.deletionVector != null) {
+          numLogicalRecordsAddedInFilesWithDeletionVectorsPartial +=
+            a.numLogicalRecords.getOrElse(0L)
+        }
       case r: RemoveFile =>
         numFilesRemoved += 1
         numLogicalRecordsRemovedPartial += r.numLogicalRecords.getOrElse {
@@ -85,6 +99,9 @@ object NumRecordsStats {
       numDeletionVectorRecordsAdded = numDeletionVectorRecordsAdded,
       numDeletionVectorRecordsRemoved = numDeletionVectorRecordsRemoved,
       numFilesAddedWithoutNumRecords = numFilesAddedWithoutNumRecords,
-      numFilesRemovedWithoutNumRecords = numFilesRemovedWithoutNumRecords)
+      numFilesRemovedWithoutNumRecords = numFilesRemovedWithoutNumRecords,
+      numLogicalRecordsAddedInFilesWithDeletionVectorsPartial =
+        numLogicalRecordsAddedInFilesWithDeletionVectorsPartial
+    )
   }
 }

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

@@ -18,6 +18,8 @@ package org.apache.spark.sql.delta.commands
 
 import java.util.concurrent.TimeUnit
 
+import scala.util.control.NonFatal
+
 import org.apache.spark.sql.delta.metric.IncrementMetric
 import org.apache.spark.sql.delta._
 import org.apache.spark.sql.delta.ClassicColumnConversions._
@@ -36,6 +38,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference,
 import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.logical.{DeltaDelete, LogicalPlan}
+import org.apache.spark.sql.delta.DeltaOperations.Operation
 import org.apache.spark.sql.execution.command.LeafRunnableCommand
 import org.apache.spark.sql.execution.metric.SQLMetric
 import org.apache.spark.sql.execution.metric.SQLMetrics.{createMetric, createTimingMetric}
@@ -127,9 +130,12 @@ case class DeleteCommand(
         }
 
         val (deleteActions, deleteMetrics) = performDelete(sparkSession, deltaLog, txn)
+        val numRecordsStats = NumRecordsStats.fromActions(deleteActions)
+        val operation = DeltaOperations.Delete(condition.toSeq)
+        validateNumRecords(deleteActions, numRecordsStats, operation)
         val commitVersion = txn.commitIfNeeded(
           actions = deleteActions,
-          op = DeltaOperations.Delete(condition.toSeq),
+          op = operation,
           tags = RowTracking.addPreservedRowTrackingTagIfNotSet(txn.snapshot))
         recordDeltaEvent(
           deltaLog,
@@ -472,6 +478,111 @@ case class DeleteCommand(
     spark.conf.get(DeltaSQLConf.DELETE_USE_PERSISTENT_DELETION_VECTORS) &&
       DeletionVectorUtils.deletionVectorsWritable(txn.snapshot)
   }
+
+  /**
+   * Validates that the number of records does not increase.
+   *
+   * Note: ideally we would also compare the number of added/removed rows in the statistics with the
+   * number of deleted/copied rows in the SQL metrics, but unfortunately this is not possible, as
+   * sql metrics are not reliable when there are task or stage retries.
+   */
+  private def validateNumRecords(
+      actions: Seq[Action],
+      numRecordsStats: NumRecordsStats,
+      op: Operation): Unit = {
+    (numRecordsStats.numLogicalRecordsAdded,
+      numRecordsStats.numLogicalRecordsRemoved,
+      numRecordsStats.numLogicalRecordsAddedInFilesWithDeletionVectors) match {
+      case (
+        Some(numAddedRecords),
+        Some(numRemovedRecords),
+        Some(numRecordsNotCopied)) =>
+        if (numAddedRecords > numRemovedRecords) {
+          logNumRecordsMismatch(deltaLog, actions, numRecordsStats, op)
+          if (conf.getConf(DeltaSQLConf.NUM_RECORDS_VALIDATION_ENABLED)) {
+            throw DeltaErrors.numRecordsMismatch(
+              operation = "DELETE",
+              numAddedRecords,
+              numRemovedRecords
+            )
+          }
+        }
+
+        if (conf.getConf(DeltaSQLConf.COMMAND_INVARIANT_CHECKS_USE_UNRELIABLE)) {
+          // and also using regular (unreliable) metrics for baseline
+          validateMetricBasedCommandInvariants(
+            numAddedRecords, numRemovedRecords, numRecordsNotCopied, op, deltaLog)
+        }
+
+      case _ =>
+        recordDeltaEvent(deltaLog, opType = "delta.assertions.statsNotPresentForNumRecordsCheck")
+        logWarning(log"Could not validate number of records due to missing statistics.")
+    }
+  }
+
+  private def validateMetricBasedCommandInvariants(
+      numAddedRecords: Long,
+      numRemovedRecords: Long,
+      numRecordsNotCopied: Long,
+      op: Operation,
+      deltaLog: DeltaLog): Unit = try {
+
+    val numRowsDeleted = CommandInvariantMetricValueFromSingle(metrics("numDeletedRows"))
+    val numRowsCopied = CommandInvariantMetricValueFromSingle(metrics("numCopiedRows"))
+
+    val recordMetricsFromMetadata = conf.getConf(DeltaSQLConf.DELTA_DML_METRICS_FROM_METADATA)
+    if (numRowsDeleted.getOrDummy == 0 && !recordMetricsFromMetadata) {
+      // If we don't record metrics we can't use them to perform invariant checks.
+      return
+    }
+
+    checkCommandInvariant(
+      invariant = () =>
+        numRowsDeleted.getOrThrow + numRowsCopied.getOrThrow + numRecordsNotCopied
+          == numRemovedRecords,
+      label = "numRowsDeleted + numRowsCopied + numRecordsNotCopied + " +
+        "numRowsRemovedByMetadataOnlyDelete == numRemovedRecords",
+      op = op,
+      deltaLog = deltaLog,
+      parameters = Map(
+        "numRowsDeleted" -> numRowsDeleted.getOrDummy,
+        "numRowsCopied" -> numRowsCopied.getOrDummy,
+        "numRemovedRecords" -> numRemovedRecords,
+        "numRecordsNotCopied" -> numRecordsNotCopied
+      ),
+      additionalInfo = Map(
+        DeltaSQLConf.DELTA_DML_METRICS_FROM_METADATA.key -> recordMetricsFromMetadata.toString
+      )
+    )
+
+    checkCommandInvariant(
+      invariant = () => numRowsCopied.getOrThrow + numRecordsNotCopied == numAddedRecords,
+      label = "numRowsCopied + numRecordsNotCopied == numAddedRecords",
+      op = op,
+      deltaLog = deltaLog,
+      parameters = Map(
+        "numRowsCopied" -> numRowsCopied.getOrDummy,
+        "numAddedRecords" -> numAddedRecords,
+        "numRecordsNotCopied" -> numRecordsNotCopied
+      ),
+      additionalInfo = Map(
+        DeltaSQLConf.DELTA_DML_METRICS_FROM_METADATA.key -> recordMetricsFromMetadata.toString
+      )
+    )
+  } catch {
+    // Immediately re-throw actual command invariant violations, so we don't re-wrap them below.
+    case e: DeltaIllegalStateException if e.getErrorClass == "DELTA_COMMAND_INVARIANT_VIOLATION" =>
+      throw e
+    case NonFatal(e) =>
+      logWarning(log"Unexpected error in validateMetricBasedCommandInvariants", e)
+      checkCommandInvariant(
+        invariant = () => false,
+        label = "Unexpected error in validateMetricBasedCommandInvariants",
+        op = op,
+        deltaLog = deltaLog,
+        parameters = Map.empty
+      )
+  }
 }
 
 object DeleteCommand {

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

@@ -21,7 +21,7 @@ import java.util.concurrent.TimeUnit.NANOSECONDS
 
 import scala.util.control.NonFatal
 
-import org.apache.spark.sql.delta.{DeltaAnalysisException, DeltaErrors, DeltaLog, DeltaOptions, DeltaTableIdentifier, DeltaTableUtils, OptimisticTransaction, ResolvedPathBasedNonDeltaTable}
+import org.apache.spark.sql.delta.{DeltaAnalysisException, DeltaErrors, DeltaLog, DeltaOptions, DeltaTableIdentifier, DeltaTableUtils, NumRecordsStats, OptimisticTransaction, ResolvedPathBasedNonDeltaTable}
 import org.apache.spark.sql.delta.actions._
 import org.apache.spark.sql.delta.catalog.{DeltaTableV2, IcebergTablePlaceHolder}
 import org.apache.spark.sql.delta.files.TahoeBatchFileIndex
@@ -40,6 +40,8 @@ import org.apache.spark.sql.catalyst.expressions.{Expression, SubqueryExpression
 import org.apache.spark.sql.catalyst.parser.ParseException
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.connector.catalog.V1Table
+import org.apache.spark.sql.delta.DeltaOperations.Operation
+import org.apache.spark.sql.delta.sources.DeltaSQLConf.DELTA_COLLECT_STATS
 import org.apache.spark.sql.execution.SQLExecution
 import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelationWithTable}
 import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
@@ -48,7 +50,7 @@ import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 /**
  * Helper trait for all delta commands.
  */
-trait DeltaCommand extends DeltaLogging {
+trait DeltaCommand extends DeltaLogging with DeltaCommandInvariants {
   /**
    * Converts string predicates into [[Expression]]s relative to a transaction.
    *
@@ -440,4 +442,50 @@ trait DeltaCommand extends DeltaLogging {
     (txnVersion, txnAppId, fromSessionConf)
   }
 
+  protected def logNumRecordsMismatch(
+      deltaLog: DeltaLog,
+      actions: Seq[Action],
+      stats: NumRecordsStats,
+      op: Operation): Unit = {
+
+    var numRemove = 0
+    var numAdd = 0
+    actions.foreach {
+      case _: AddFile =>
+        numAdd += 1
+      case _: RemoveFile =>
+        numRemove += 1
+      case _ =>
+    }
+
+    val info = NumRecordsCheckInfo(
+      operation = op.name,
+      numAdd = numAdd,
+      numRemove = numRemove,
+      numRecordsRemoved = stats.numLogicalRecordsRemovedPartial,
+      numRecordsAdded = stats.numLogicalRecordsAddedPartial,
+      numDeletionVectorRecordsRemoved = stats.numDeletionVectorRecordsRemoved,
+      numDeletionVectorRecords = stats.numDeletionVectorRecordsAdded,
+      operationParameters = op.jsonEncodedValues,
+      statsCollectionEnabled = SparkSession.getActiveSession.get.conf.get(DELTA_COLLECT_STATS)
+    )
+    recordDeltaEvent(deltaLog, opType = "delta.assertions.numRecordsChanged", data = info)
+    logWarning(log"Number of records validation failed. Number of added records" +
+      log" (${MDC(DeltaLogKeys.NUM_RECORDS, stats.numLogicalRecordsAddedPartial)})" +
+      log" does not match removed records" +
+      log" (${MDC(DeltaLogKeys.NUM_RECORDS2, stats.numLogicalRecordsRemovedPartial)})")
+  }
 }
+
+// Recorded when number of records check for unchanged data fails.
+case class NumRecordsCheckInfo(
+    operation: String,
+    numAdd: Int,
+    numRemove: Int,
+    numRecordsAdded: Long,
+    numRecordsRemoved: Long,
+    numDeletionVectorRecordsRemoved: Long = 0, // number of DV records removed by the RemoveFiles
+    numDeletionVectorRecords: Long = 0, // number of DV records present in all AddFiles
+    operationParameters: Map[String, String],
+    statsCollectionEnabled: Boolean
+)