[Spark] Materialize cached Delta sources in MERGE (delta-io#4629)

## Description


Fixes a source of non-determinism in MERGE due to query caching. Query
caching doesn't pin the version of Delta tables, a source cached using
`df.cache()` can then return different results between the two internal
MERGE jobs if it is updated in the meantime.

Guarded by flag `spark.databricks.delta.merge.materializeCachedSource`.

## How was this patch tested?
This is tough to cover with a test since it requires precise timing to
update the source while the MERGE operation is running.
I did validate that the fix works manually, running with a debugger.

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

Author: johanl-db

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

@@ -1689,6 +1689,12 @@
     ],
     "sqlState" : "42601"
   },
+  "DELTA_MERGE_SOURCE_CACHED_DURING_EXECUTION" : {
+    "message" : [
+      "The MERGE operation failed because (part of) the source plan was cached while the MERGE operation was running."
+    ],
+    "sqlState" : "25000"
+  },
   "DELTA_MERGE_UNEXPECTED_ASSIGNMENT_KEY" : {
     "message" : [
       "Unexpected assignment key: <unexpectedKeyClass> - <unexpectedKeyObject>"

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

@@ -1325,6 +1325,9 @@ trait DeltaErrorsBase
     )
   }
 
+  def mergeConcurrentOperationCachedSourceException(): Throwable =
+    new DeltaRuntimeException(errorClass = "DELTA_MERGE_SOURCE_CACHED_DURING_EXECUTION")
+
   def columnOfTargetTableNotFoundInMergeException(targetCol: String,
       colNames: String): Throwable = {
     new DeltaAnalysisException(

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

@@ -481,6 +481,27 @@ trait MergeIntoCommandBase extends LeafRunnableCommand
     }
   }
 
+  /**
+   * Check whether (part of) the give source dataframe is cached and logs an assertion or fails if
+   * it is. Query caching doesn't pin versions of delta tables and can lead to incorrect results so
+   * cached source plans must be materialized.
+   */
+  def checkSourcePlanIsNotCached(spark: SparkSession, source: LogicalPlan): Unit = {
+    val sourceIsCached = planContainsCachedRelation(DataFrameUtils.ofRows(spark, source))
+    if (sourceIsCached &&
+        spark.conf.get(DeltaSQLConf.MERGE_FAIL_SOURCE_CACHED_AFTER_MATERIALIZATION)) {
+      throw DeltaErrors.mergeConcurrentOperationCachedSourceException()
+    }
+
+    deltaAssert(
+      !sourceIsCached,
+      name = "merge.sourceCachedAfterMaterializationStep",
+      msg = "Cached source plans must be materialized in MERGE but the source only got cached " +
+        "after the decision to materialize was taken.",
+      deltaLog = targetDeltaLog
+    )
+  }
+
   override protected def prepareMergeSource(
       spark: SparkSession,
       source: LogicalPlan,

spark/src/main/scala/org/apache/spark/sql/delta/commands/merge/ClassicMergeExecutor.scala

@@ -165,6 +165,8 @@ trait ClassicMergeExecutor extends MergeOutputGeneration {
       .collect()
       .head
 
+    checkSourcePlanIsNotCached(spark, getMergeSource.df.queryExecution.logical)
+
     val hasMultipleMatches = multipleMatchCount > 0
     throwErrorOnMultipleMatches(hasMultipleMatches, spark)
     if (hasMultipleMatches) {
@@ -465,6 +467,8 @@ trait ClassicMergeExecutor extends MergeOutputGeneration {
     // Write to Delta
     val newFiles = writeFiles(spark, deltaTxn, outputDF)
 
+    checkSourcePlanIsNotCached(spark, getMergeSource.df.queryExecution.logical)
+
     // Update metrics
     val (addedBytes, addedPartitions) = totalBytesAndDistinctPartitionValues(newFiles)
     metrics("numTargetFilesAdded") += newFiles.count(_.isInstanceOf[AddFile])

spark/src/main/scala/org/apache/spark/sql/delta/commands/merge/MergeIntoMaterializeSource.scala

@@ -246,6 +246,8 @@ trait MergeIntoMaterializeSource extends DeltaLogging with DeltaSparkPlanUtils {
     import DeltaSQLConf.MergeMaterializeSource._
     val checkDeterministicOptions =
       DeltaSparkPlanUtils.CheckDeterministicOptions(allowDeterministicUdf = true)
+    val materializeCachedSource = spark.conf.get(DeltaSQLConf.MERGE_MATERIALIZE_CACHED_SOURCE)
+
     materializeType match {
       case ALL =>
         (true, MergeIntoMaterializeSourceReason.MATERIALIZE_ALL)
@@ -270,6 +272,12 @@ trait MergeIntoMaterializeSource extends DeltaLogging with DeltaSparkPlanUtils {
           // the user defined function is marked as deterministic, as it is often incorrectly marked
           // as such.
           (true, MergeIntoMaterializeSourceReason.NON_DETERMINISTIC_SOURCE_WITH_DETERMINISTIC_UDF)
+        } else if (materializeCachedSource &&
+            planContainsCachedRelation(DataFrameUtils.ofRows(spark, source))) {
+          // The query cache doesn't pin the version of cached Delta tables, cache can get
+          // concurrently updated in the middle of MERGE execution. We materialize the source in
+          // that case to avoid this issue.
+          (true, MergeIntoMaterializeSourceReason.SOURCE_CACHED)
         } else {
           (false, MergeIntoMaterializeSourceReason.NOT_MATERIALIZED_AUTO)
         }
@@ -460,6 +468,8 @@ object MergeIntoMaterializeSourceReason extends Enumeration {
     Value("materializeNonDeterministicSourceWithDeterministicUdf")
   // Materialize when the configuration is invalid
   val INVALID_CONFIG = Value("invalidConfigurationFailsafe")
+  // Materialize when the source is cached.
+  val SOURCE_CACHED = Value("materializeCachedSource")
   // Catch-all case.
   val UNKNOWN = Value("unknown")
 
@@ -470,7 +480,8 @@ object MergeIntoMaterializeSourceReason extends Enumeration {
     NON_DETERMINISTIC_SOURCE_OPERATORS,
     IGNORE_UNREADABLE_FILES_CONFIGS_ARE_SET,
     NON_DETERMINISTIC_SOURCE_WITH_DETERMINISTIC_UDF,
-    INVALID_CONFIG
+    INVALID_CONFIG,
+    SOURCE_CACHED
   )
 }
 

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

@@ -861,6 +861,32 @@ trait DeltaSQLConfBase {
       .booleanConf
       .createWithDefault(true)
 
+  val MERGE_MATERIALIZE_CACHED_SOURCE =
+    buildConf("merge.materializeCachedSource")
+      .internal()
+      .doc(
+        """
+          |When enabled, materialize the source in MERGE if it is cached (e.g. via df.cache()). This
+          |prevents incorrect results due to query caching not pinning the version of cached Delta
+          |tables.
+          |""".stripMargin)
+      .booleanConf
+      .createWithDefault(true)
+
+  val MERGE_FAIL_SOURCE_CACHED_AFTER_MATERIALIZATION =
+    buildConf("merge.failSourceCachedAfterMaterialization")
+      .internal()
+      .doc(
+        """
+          |Enables a check that fails the MERGE operation if the source was cached (using
+          |df.cache()) after the source materialization phase. Query caching doesn't pin the version
+          |of Delta tables and we should materialize cached source plans. In rare cases, the source
+          |might get cached after the decision to materialize, which could lead to incorrect results
+          |if we let the operation succeed.
+          |""".stripMargin)
+      .booleanConf
+      .createWithDefault(true)
+
   val MERGE_MATERIALIZE_SOURCE_RDD_STORAGE_LEVEL =
     buildConf("merge.materializeSource.rddStorageLevel")
       .internal()

spark/src/main/scala/org/apache/spark/sql/delta/util/DeltaSparkPlanUtils.scala

@@ -18,8 +18,10 @@ package org.apache.spark.sql.delta.util
 
 import org.apache.spark.sql.delta.{DeltaTable, DeltaTableReadPredicate}
 
+import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.catalyst.expressions.{Exists, Expression, InSubquery, LateralSubquery, ScalarSubquery, SubqueryExpression => SparkSubqueryExpression, UserDefinedExpression}
 import org.apache.spark.sql.catalyst.plans.logical.{Distinct, Filter, LeafNode, LogicalPlan, OneRowRelation, Project, SubqueryAlias, Union}
+import org.apache.spark.sql.execution.columnar.InMemoryRelation
 import org.apache.spark.sql.execution.datasources.LogicalRelation
 
 
@@ -45,6 +47,10 @@ trait DeltaSparkPlanUtils {
     )
   }
 
+  /** Returns whether part of the plan was cached using df.cache() or similar. */
+  protected def planContainsCachedRelation(df: DataFrame): Boolean =
+    df.queryExecution.withCachedData.exists(_.isInstanceOf[InMemoryRelation])
+
   /**
    * Returns `true` if `plan` has a safe level of determinism. This is a conservative
    * approximation of `plan` being a truly deterministic query.