[RLC] Case 1 follow-up — MERGE source-key read predicate as an exact semi-join (MERGE ‖ MERGE on disjoint keys)

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

@@ -23,21 +23,25 @@ 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._
 import org.apache.spark.sql.delta.actions.{Action, AddFile, FileAction}
+import org.apache.spark.sql.delta.expressions.MergeSourceKeyMatch
 import org.apache.spark.sql.delta.commands.merge.{MergeIntoMaterializeSource, MergeIntoMaterializeSourceReason, MergeStats}
 import org.apache.spark.sql.delta.files.{TahoeBatchFileIndex, TahoeFileIndex, TransactionalWrite}
 import org.apache.spark.sql.delta.metering.DeltaLogging
 import org.apache.spark.sql.delta.schema.{ImplicitMetadataOperation, SchemaUtils}
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
 import org.apache.spark.SparkContext
-import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SparkSession}
+import org.apache.spark.sql.{AnalysisException, Column, DataFrame, Row, SparkSession}
+import org.apache.spark.sql.catalyst.CatalystTypeConverters
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
 import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
 import org.apache.spark.sql.execution.command.LeafRunnableCommand
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
-import org.apache.spark.sql.functions.col
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.functions.{col, max, min}
+import org.apache.spark.sql.types.{AtomicType, DataType, DoubleType, FloatType, StringType, StructType}
 
 trait MergeIntoCommandBase extends LeafRunnableCommand
   with DeltaCommand
@@ -383,6 +387,114 @@ trait MergeIntoCommandBase extends LeafRunnableCommand
 
   protected def seqToString(exprs: Seq[Expression]): String = exprs.map(_.sql).mkString("\n\t")
 
+  /**
+   * The `(target, source)` expression pairs of the `target_expr = source_expr` conjuncts of the ON
+   * condition, where each side is deterministic and references only its own side, and both sides
+   * have the same atomic type whose equality is binary (so a pre-filter on the collected source
+   * values agrees with the join's `=`). Excluded: null-safe equality (NULL keys would match),
+   * FLOAT/DOUBLE (`-0.0 = 0.0` and NaN normalization) and non-binary string collations.
+   */
+  private def sourceTargetEquiJoinKeys: Seq[(Expression, Expression)] = {
+    def binaryComparable(dt: DataType): Boolean = dt match {
+      case _: FloatType | _: DoubleType => false
+      case st: StringType => st.supportsBinaryOrdering
+      case _: AtomicType => true
+      case _ => false
+    }
+    def onlyReferences(e: Expression, side: AttributeSet): Boolean =
+      e.deterministic && e.references.nonEmpty && e.references.subsetOf(side) &&
+        binaryComparable(e.dataType)
+    def keyPair(t: Expression, s: Expression): Option[(Expression, Expression)] =
+      if (onlyReferences(t, target.outputSet) && onlyReferences(s, source.outputSet) &&
+          DataTypeUtils.sameType(t.dataType, s.dataType)) {
+        Some((t, s))
+      } else {
+        None
+      }
+    splitConjunctivePredicates(condition).flatMap {
+      case EqualTo(l, r) => keyPair(l, r).orElse(keyPair(r, l))
+      case _ => None
+    }
+  }
+
+  /**
+   * Read predicates recording which target rows this MERGE depends on: those whose equi-join
+   * key tuple equals the key tuple of some source row. Mirrors DBR's MERGE source read predicate
+   * (a join against the MERGE source, with no limit on the number of source keys).
+   *
+   * A MERGE reads exactly the target rows that join with some source row (plus, for NOT MATCHED
+   * BY SOURCE, every other target row, so such MERGEs are excluded). Any such row's key tuple is
+   * in the source's key tuples, so "key tuple in source keys" is a sound over-approximation of
+   * the rows read, including the rows whose absence lets a NOT MATCHED clause insert, and the
+   * rows a MATCHED clause condition only narrows.
+   *
+   * Returned predicates (all ANDed in one read):
+   *  - per key, a cheap pre-filter on the source values: `target_key IN (values)` when there are at
+   *    most [[DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_IN_SET_THRESHOLD]] distinct key tuples,
+   *    else `target_key BETWEEN min AND max`. It prunes this MERGE's own target scan and is what
+   *    file-level (stats and partition) conflict checks use. A key with no non-null source value
+   *    matches no target row, so it becomes `false`.
+   *  - one [[MergeSourceKeyMatch]] carrying the source keys; the value-exact conflict checks
+   *    evaluate it as an exact semi-join, and it is `true` everywhere else.
+   *
+   * The source is the prepared (materialized, or deterministic) source the MERGE joins with, so the
+   * keys seen here and by the conflict check are the keys the MERGE join sees.
+   */
+  protected def getSourceKeyReadPredicates(spark: SparkSession): Seq[Expression] = {
+    if (!spark.conf.get(DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_ENABLED) ||
+        notMatchedBySourceClauses.nonEmpty) {
+      return Seq.empty
+    }
+    val keyPairs = sourceTargetEquiJoinKeys
+    if (keyPairs.isEmpty) {
+      return Seq.empty
+    }
+    val inSetThreshold =
+      spark.conf.get(DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_IN_SET_THRESHOLD)
+
+    val keyNames = keyPairs.indices.map(i => s"__merge_source_key_$i")
+    val sourceKeys = getMergeSource.df.select(keyPairs.zip(keyNames).map {
+      case ((_, sourceKey), name) => Column(Alias(sourceKey, name)())
+    }: _*)
+    val keyRows =
+      if (inSetThreshold > 0) sourceKeys.distinct().limit(inSetThreshold + 1).collect()
+      else Array.empty[Row]
+    val useInSet = inSetThreshold > 0 && keyRows.length <= inSetThreshold
+
+    val preFilters: Seq[Expression] = if (useInSet) {
+      keyPairs.zipWithIndex.map { case ((targetKey, sourceKey), i) =>
+        val toCatalyst = CatalystTypeConverters.createToCatalystConverter(sourceKey.dataType)
+        val values = keyRows.iterator.map(row => toCatalyst(row.get(i))).filter(_ != null).toSet
+        if (values.isEmpty) Literal.FalseLiteral else InSet(targetKey, values)
+      }
+    } else {
+      val bounds = keyNames.flatMap(n => Seq(min(col(n)), max(col(n)))) match {
+        case aggs => sourceKeys.agg(aggs.head, aggs.tail: _*).head()
+      }
+      keyPairs.zipWithIndex.map { case ((targetKey, sourceKey), i) =>
+        val toCatalyst = CatalystTypeConverters.createToCatalystConverter(sourceKey.dataType)
+        val (lo, hi) = (toCatalyst(bounds.get(2 * i)), toCatalyst(bounds.get(2 * i + 1)))
+        if (lo == null) {
+          Literal.FalseLiteral
+        } else {
+          And(
+            GreaterThanOrEqual(targetKey, Literal(lo, sourceKey.dataType)),
+            LessThanOrEqual(targetKey, Literal(hi, sourceKey.dataType)))
+        }
+      }
+    }
+
+    recordDeltaEvent(
+      targetDeltaLog,
+      opType = "delta.dml.merge.sourceKeyReadPredicate",
+      data = Map(
+        "numKeyColumns" -> keyPairs.size,
+        "inSetThreshold" -> inSetThreshold,
+        "preFilter" -> (if (useInSet) "inSet" else "range"),
+        "numDistinctKeys" -> (if (useInSet) keyRows.length.toLong else -1L)))
+    preFilters :+ MergeSourceKeyMatch(keyPairs.map(_._1), sourceKeys)
+  }
+
   /**
    * Execute the given `thunk` and return its result while recording the time taken to do it
    * and setting additional local properties for better UI visibility.

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

@@ -84,9 +84,13 @@ trait ClassicMergeExecutor extends MergeOutputGeneration {
     spark.sparkContext.register(touchedFilesAccum, TOUCHED_FILES_ACCUM_NAME)
 
     // Prune non-matching files if we don't need to collect them for NOT MATCHED BY SOURCE clauses.
+    // The optional source-key predicates also record which target keys this MERGE reads, so that
+    // conflict detection can tell MERGEs on disjoint keys apart.
     val dataSkippedFiles =
       if (notMatchedBySourceClauses.isEmpty) {
-        deltaTxn.filterFiles(getTargetOnlyPredicates(spark), keepNumRecords = true)
+        deltaTxn.filterFiles(
+          getTargetOnlyPredicates(spark) ++ getSourceKeyReadPredicates(spark),
+          keepNumRecords = true)
       } else {
         deltaTxn.filterFiles(filters = Seq(Literal.TrueLiteral), keepNumRecords = true)
       }

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

@@ -84,7 +84,8 @@ trait InsertOnlyMergeExecutor extends MergeOutputGeneration {
         val conjunctivePredicates = splitConjunctivePredicates(condition)
         val targetOnlyPredicates =
           conjunctivePredicates.filter(_.references.subsetOf(target.outputSet))
-        dataSkippedFiles = Some(deltaTxn.filterFiles(targetOnlyPredicates))
+        dataSkippedFiles = Some(
+          deltaTxn.filterFiles(targetOnlyPredicates ++ getSourceKeyReadPredicates(spark)))
 
         val targetPlan = buildTargetPlanWithFiles(
           spark,

spark/src/main/scala/org/apache/spark/sql/delta/expressions/MergeSourceKeyMatch.scala

@@ -0,0 +1,62 @@
+/*
+ * Copyright (2021) The Delta Lake Project Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.delta.expressions
+
+import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Expression, Predicate}
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, EmptyBlock, ExprCode, FalseLiteral, TrueLiteral}
+
+/**
+ * A MERGE read predicate meaning "the target row's key tuple `targetKeys` equals the key tuple of
+ * some row of `sourceKeys`" -- the rows a MERGE whose ON condition has the equi-join conjuncts
+ * `targetKeys(i) = source_key_i` can read.
+ *
+ * Conflict detection evaluates it exactly, as a left-semi join of the concurrently changed rows
+ * against `sourceKeys` (see `ConflictDataSkippingReader`); this mirrors DBR's
+ * `DeltaTableMergeSourceReadPredicate`, so there is no limit on the number of source keys.
+ * Everywhere else (file pruning, partition filters, stats skipping) it evaluates to `true`: it is
+ * always recorded next to cheap over-approximating pre-filters on the same keys, which do the
+ * pruning, so treating it as `true` there only keeps more rows, never fewer.
+ *
+ * @param targetKeys the target side of each equi-join conjunct (its only children, so attribute
+ *                   rebinding / rewriting applies to them).
+ * @param sourceKeys the MERGE source projected to the source side of each conjunct, column `i`
+ *                   aligned with `targetKeys(i)`. Driver-only (never shipped to executors).
+ */
+case class MergeSourceKeyMatch(
+    targetKeys: Seq[Expression],
+    @transient sourceKeys: DataFrame)
+  extends Expression with Predicate {
+
+  override def children: Seq[Expression] = targetKeys
+
+  override def nullable: Boolean = false
+
+  override def eval(input: InternalRow): Any = true
+
+  override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode =
+    ev.copy(code = EmptyBlock, isNull = FalseLiteral, value = TrueLiteral)
+
+  override protected def withNewChildrenInternal(
+      newChildren: IndexedSeq[Expression]): MergeSourceKeyMatch =
+    copy(targetKeys = newChildren)
+
+  override def toString: String = s"merge_source_key_match(${targetKeys.mkString(", ")})"
+
+  override def sql: String = s"merge_source_key_match(${targetKeys.map(_.sql).mkString(", ")})"
+}

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

@@ -586,6 +586,39 @@ trait DeltaSQLConfBase extends DeltaSQLConfUtils {
       .longConf
       .createWithDefault(1L << 30) // 1 GiB
 
+  val MERGE_SOURCE_KEY_READ_PREDICATE_ENABLED =
+    buildConf("merge.sourceKeyReadPredicate.enabled")
+      .internal()
+      .doc(
+        """When enabled, a MERGE without NOT MATCHED BY SOURCE clauses records, as its read
+          |predicate, the target rows whose `target_expr = source_expr` equi-join keys (from the ON
+          |condition) equal the keys of some source row. Without it, a key-only ON condition such as
+          |`t.id = s.id` records a read predicate of `true`, so any concurrently added file
+          |conflicts with the MERGE. The predicate is a cheap pre-filter on the source keys (see
+          |merge.sourceKeyReadPredicate.inSetThreshold), which also prunes the MERGE's own target
+          |scan, plus an exact match that the value-exact conflict checks
+          |(conflictDetection.dataSkipping.valueExact.enabled) evaluate as a semi-join against the
+          |MERGE source, like Databricks Runtime. Concurrent MERGEs on disjoint keys then commit,
+          |while a concurrent change to any of this MERGE's keys (an update, or an insert of the
+          |same new key) still conflicts. There is no limit on the number of source keys.
+          |""".stripMargin)
+      .booleanConf
+      .createWithDefault(false)
+
+  val MERGE_SOURCE_KEY_READ_PREDICATE_IN_SET_THRESHOLD =
+    buildConf("merge.sourceKeyReadPredicate.inSetThreshold")
+      .internal()
+      .doc(
+        """The maximum number of distinct source key tuples for which
+          |merge.sourceKeyReadPredicate.enabled builds a `target_expr IN (source values)`
+          |pre-filter (the values are collected to the driver). Above it, the pre-filter is a
+          |`target_expr BETWEEN min AND max` range of the source values. This only affects pruning:
+          |the exact source-key match still decides the conflict, so a smaller value never causes
+          |a conflict the value-exact check would not. A non-positive value always uses the
+          |range.""".stripMargin)
+      .intConf
+      .createWithDefault(10000)
+
   val DELTA_PROTOCOL_DEFAULT_WRITER_VERSION =
     buildConf("properties.defaults.minWriterVersion")
       .doc("The default writer protocol version to create new tables with, unless a feature " +

spark/src/main/scala/org/apache/spark/sql/delta/stats/ConflictDataSkippingReader.scala

@@ -21,7 +21,7 @@ import scala.util.control.NonFatal
 import org.apache.spark.sql.delta.ClassicColumnConversions._
 import org.apache.spark.sql.delta.DeltaTableUtils
 import org.apache.spark.sql.delta.actions.AddFile
-import org.apache.spark.sql.delta.expressions.DecodeNestedZ85EncodedVariant
+import org.apache.spark.sql.delta.expressions.{DecodeNestedZ85EncodedVariant, MergeSourceKeyMatch}
 import org.apache.spark.sql.delta.metering.DeltaLogging
 import org.apache.spark.sql.delta.schema.SchemaUtils
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
@@ -270,19 +270,14 @@ trait ConflictDataSkippingReader extends DeltaLogging { self: DataSkippingReader
       // winner's files and this schema agree by the time we get here.
       val df = snapshot.deltaLog.createDataFrame(
         snapshot, files, actionTypeOpt = Some(ConflictDataSkippingReader.VALUE_EXACT_ACTION_TYPE))
-      // Rebind each filter's attributes to the fresh DataFrame's columns by name, AND within a
-      // read, OR across reads -- a file matches if any read's predicate matches any of its rows.
-      val condition = perReadEligible
-        .map(filters => filters.map(f => rebindToDataFrame(f, df)).reduce(_ && _))
-        .reduce(_ || _)
       // The caller only needs a boolean (see the scaladoc), so run one short-circuiting existence
       // check instead of attributing rows back to files.
       val sc = spark.sparkContext
       val prevJobDesc = sc.getLocalProperty(SparkContext.SPARK_JOB_DESCRIPTION)
       try {
         sc.setJobDescription("Delta conflict detection: value-exact data skipping")
         recordFrameProfile("Delta", "DataSkippingReader.anyAddedRowMatchesReadPredicate") {
-          !df.where(condition).isEmpty
+          !rowsMatchingAnyRead(df, perReadEligible).isEmpty
         }
       } finally {
         sc.setJobDescription(prevJobDesc)
@@ -348,10 +343,7 @@ trait ConflictDataSkippingReader extends DeltaLogging { self: DataSkippingReader
               val df = snapshot.deltaLog.createDataFrame(
                 snapshot, files,
                 actionTypeOpt = Some(ConflictDataSkippingReader.DELETE_READ_ACTION_TYPE))
-              val condition = perReadEligible
-                .map(filters => filters.map(f => rebindToDataFrame(f, df)).reduce(_ && _))
-                .reduce(_ || _)
-              df.where(condition).select(lit(weight).as("weight"))
+              rowsMatchingAnyRead(df, perReadEligible).select(lit(weight).as("weight"))
             }
             // matches(removed) = matches(pre-image live) - matches(post-image live), evaluated in a
             // SINGLE Spark job: union the per-image weighted matches (+1 pre, -1 post) and sum. A
@@ -376,6 +368,46 @@ trait ConflictDataSkippingReader extends DeltaLogging { self: DataSkippingReader
     }
   }
 
+  /**
+   * The rows of `df` matching ANY read: AND within a read, OR across reads, with every filter
+   * rebound to `df`'s columns. A [[MergeSourceKeyMatch]] is evaluated exactly, as a left-semi join
+   * of the read's (already filtered) rows against its source keys on the key tuple -- the same join
+   * DBR uses for a MERGE source read predicate -- so a MERGE's read covers just the keys of its
+   * source, however many there are. Without one, this is a single `where`.
+   *
+   * With a join, the reads are unioned, so a row matching several reads appears once per read.
+   * That multiplicity depends only on the row's values, so existence checks and the delete/read
+   * pre-minus-post count (where each live post-image row is also a pre-image row) stay exact.
+   */
+  private def rowsMatchingAnyRead(
+      df: DataFrame,
+      perReadEligible: Seq[Seq[Expression]]): DataFrame = {
+    def isKeyMatch(e: Expression): Boolean = e.isInstanceOf[MergeSourceKeyMatch]
+    def where(rows: DataFrame, filters: Seq[Expression]): DataFrame =
+      filters.map(f => rebindToDataFrame(f, df)).reduceOption(_ && _)
+        .map(condition => rows.where(condition))
+        .getOrElse(rows)
+    if (!perReadEligible.exists(_.exists(isKeyMatch))) {
+      df.where(perReadEligible
+        .map(filters => filters.map(f => rebindToDataFrame(f, df)).reduce(_ && _))
+        .reduce(_ || _))
+    } else {
+      perReadEligible.map { filters =>
+        val (keyMatches, plainFilters) = filters.partition(isKeyMatch)
+        keyMatches.foldLeft(where(df, plainFilters)) {
+          case (rows, m: MergeSourceKeyMatch) =>
+            val source = m.sourceKeys
+            val joinCondition = m.targetKeys.zip(source.columns).map {
+              case (targetKey, sourceKey) =>
+                rebindToDataFrame(targetKey, df) === source.col(quoteColumnName(sourceKey))
+            }.reduce(_ && _)
+            rows.join(source, joinCondition, "left_semi")
+          case (rows, _) => rows
+        }
+      }.reduce(_ union _)
+    }
+  }
+
   /**
    * Rebinds `e`'s attribute references to `df`'s output columns BY NAME, so a read predicate
    * resolved against the transaction's plan can be evaluated on a freshly built DataFrame. Nested