From fa625807a686a5806aebdb7f6d1b1d54a025c7f7 Mon Sep 17 00:00:00 2001 From: Eunjin Song Date: Tue, 6 Oct 2026 21:08:28 -0700 Subject: [PATCH 1/2] =?UTF-8?q?[RLC]=20Case=201=20follow-up=20=E2=80=94=20?= =?UTF-8?q?MERGE=20source-key=20read=20predicate=20(MERGE=20=E2=80=96=20ME?= =?UTF-8?q?RGE=20on=20disjoint=20keys)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With a key-only ON condition (e.g. `t.id = s.id`) MERGE registers a read predicate of `true` on unpartitioned / liquid-clustered tables, so any concurrent MERGE that adds a file conflicts, even when the two MERGEs touch disjoint keys. DBR resolves this per key. This change, behind the internal flag `spark.databricks.delta.merge.sourceKeyReadPredicate.enabled` (default off), collects the distinct source values of each `target = source` equi-join conjunct (bounded by `merge.sourceKeyReadPredicate.maxKeys`, default 100000, provisional) and adds `target_key IN (source keys)` (InSet) to the predicates passed to `filterFiles`. The predicate: - is recorded as a read predicate, so concurrent appends are checked against it with min/max and value-exact conflict data skipping; - prunes the MERGE's own target scan and read-file set. Skipped for NOT MATCHED BY SOURCE (whole-table read), null-safe equality, non-atomic or mismatched key types, and when maxKeys is exceeded. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../delta/commands/MergeIntoCommandBase.scala | 84 +++++- .../commands/merge/ClassicMergeExecutor.scala | 6 +- .../merge/InsertOnlyMergeExecutor.scala | 3 +- .../sql/delta/sources/DeltaSQLConf.scala | 32 ++ .../MergeSourceKeyReadPredicateSuite.scala | 273 ++++++++++++++++++ 5 files changed, 394 insertions(+), 4 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/delta/MergeSourceKeyReadPredicateSuite.scala diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/commands/MergeIntoCommandBase.scala b/spark/src/main/scala/org/apache/spark/sql/delta/commands/MergeIntoCommandBase.scala index 5ad1b2e31f0..6847f218bd2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/delta/commands/MergeIntoCommandBase.scala +++ b/spark/src/main/scala/org/apache/spark/sql/delta/commands/MergeIntoCommandBase.scala @@ -23,6 +23,7 @@ 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.commands.merge.{MergeIntoMaterializeSource, MergeIntoMaterializeSourceReason, MergeStats} import org.apache.spark.sql.delta.files.{TahoeBatchFileIndex, TahoeFileIndex, TransactionalWrite} @@ -30,14 +31,16 @@ 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.types.{AtomicType, StructType} trait MergeIntoCommandBase extends LeafRunnableCommand with DeltaCommand @@ -383,6 +386,83 @@ 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, atomic-typed and references only its own side. + * Null-safe equality is excluded: a NULL source key would have to match NULL target keys, which + * an IN predicate cannot express. + */ + private def sourceTargetEquiJoinKeys: Seq[(Expression, Expression)] = { + def onlyReferences(e: Expression, side: AttributeSet): Boolean = + e.deterministic && e.references.nonEmpty && e.references.subsetOf(side) && + e.dataType.isInstanceOf[AtomicType] + 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 `target_expr IN (distinct source keys)`, one per equi-join key of the ON + * condition, recorded so that conflict detection knows which target rows this MERGE depends on. + * + * 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). Each target row that can satisfy + * the ON condition has every equi-join key value in the set of that key's non-null source values, + * so the conjunction of the per-key IN predicates is a sound over-approximation of the rows read, + * including the rows whose absence lets a NOT MATCHED clause insert. It is also sound for MATCHED + * clause conditions, which only narrow the joined rows. + * + * Returns no predicate (the MERGE keeps its existing read predicates) when the feature is off, + * the ON condition has no usable equi-join key, a key has no non-null source value, or there are + * more than [[DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_MAX_KEYS]] distinct key tuples. The + * source is the same prepared (materialized, or deterministic) source the MERGE joins with, so + * the collected keys are the keys the join sees. + */ + protected def getSourceKeyReadPredicates(spark: SparkSession): Seq[Expression] = { + val maxKeys = spark.conf.get(DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_MAX_KEYS) + if (!spark.conf.get(DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_ENABLED) || + maxKeys <= 0 || notMatchedBySourceClauses.nonEmpty) { + return Seq.empty + } + val keyPairs = sourceTargetEquiJoinKeys + if (keyPairs.isEmpty) { + return Seq.empty + } + + val keyColumns = keyPairs.zipWithIndex.map { case ((_, sourceKey), i) => + Column(Alias(sourceKey, s"__merge_source_key_$i")()) + } + val keyRows = getMergeSource.df.select(keyColumns: _*).distinct().limit(maxKeys + 1).collect() + val predicates = if (keyRows.length > maxKeys) { + None + } else { + val perKey = 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) None else Some(InSet(targetKey, values)) + } + if (perKey.forall(_.isDefined)) Some(perKey.flatten) else None + } + + recordDeltaEvent( + targetDeltaLog, + opType = "delta.dml.merge.sourceKeyReadPredicate", + data = Map( + "numKeyColumns" -> keyPairs.size, + "numDistinctKeys" -> keyRows.length, + "maxKeys" -> maxKeys, + "applied" -> predicates.isDefined)) + predicates.getOrElse(Seq.empty) + } + /** * 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. diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/commands/merge/ClassicMergeExecutor.scala b/spark/src/main/scala/org/apache/spark/sql/delta/commands/merge/ClassicMergeExecutor.scala index 6fe0ef8baa9..da05b8de451 100644 --- a/spark/src/main/scala/org/apache/spark/sql/delta/commands/merge/ClassicMergeExecutor.scala +++ b/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) } diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/commands/merge/InsertOnlyMergeExecutor.scala b/spark/src/main/scala/org/apache/spark/sql/delta/commands/merge/InsertOnlyMergeExecutor.scala index 4b9f6e9be3d..06da208be7f 100644 --- a/spark/src/main/scala/org/apache/spark/sql/delta/commands/merge/InsertOnlyMergeExecutor.scala +++ b/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, diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSQLConf.scala b/spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSQLConf.scala index f86fc7af694..166e4f5def7 100644 --- a/spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSQLConf.scala +++ b/spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSQLConf.scala @@ -586,6 +586,38 @@ 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, MERGE (without NOT MATCHED BY SOURCE clauses) collects the distinct + |values of the source side of each `target_expr = source_expr` equi-join conjunct of the + |ON condition and records `target_expr IN (source keys)` as an extra read predicate of + |the transaction (it is also used to prune target files for the MERGE's own scan). + |Without it, a key-only ON condition such as `t.id = s.id` registers a read predicate of + |`true`, so any concurrently added file conflicts with the MERGE. With it, conflict-time + |data skipping (conflictDetection.dataSkipping.enabled, and especially the value-exact + |tier) can prove that concurrent MERGEs on disjoint keys do not conflict, while a + |concurrent change to any of this MERGE's keys (an update, or an insert of the same new + |key) still conflicts. When the number of distinct source key tuples exceeds + |merge.sourceKeyReadPredicate.maxKeys the predicate is not recorded and MERGE behaves + |exactly as when this flag is off.""".stripMargin) + .booleanConf + .createWithDefault(false) + + val MERGE_SOURCE_KEY_READ_PREDICATE_MAX_KEYS = + buildConf("merge.sourceKeyReadPredicate.maxKeys") + .internal() + .doc( + """The maximum number of distinct source key tuples for which + |merge.sourceKeyReadPredicate.enabled records a `target_expr IN (source keys)` read + |predicate. The keys are collected to the driver, so this bounds driver memory and the + |size of the predicate evaluated during file pruning and conflict detection. Above the + |limit the MERGE falls back to its existing read predicates. A non-positive value + |disables the feature.""".stripMargin) + .intConf + .createWithDefault(100000) + val DELTA_PROTOCOL_DEFAULT_WRITER_VERSION = buildConf("properties.defaults.minWriterVersion") .doc("The default writer protocol version to create new tables with, unless a feature " + diff --git a/spark/src/test/scala/org/apache/spark/sql/delta/MergeSourceKeyReadPredicateSuite.scala b/spark/src/test/scala/org/apache/spark/sql/delta/MergeSourceKeyReadPredicateSuite.scala new file mode 100644 index 00000000000..7613a091747 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/delta/MergeSourceKeyReadPredicateSuite.scala @@ -0,0 +1,273 @@ +/* + * 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 + +import java.io.File + +import scala.concurrent.duration.Duration +import scala.util.control.NonFatal + +import com.databricks.spark.util.Log4jUsageLogger +import org.apache.spark.sql.delta.concurrency.{PhaseLockingTestMixin, TransactionExecutionTestMixin} +import org.apache.spark.sql.delta.sources.DeltaSQLConf +import org.apache.spark.sql.delta.test.DeltaSQLCommandTest +import org.apache.spark.sql.delta.util.JsonUtils + +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.util.ThreadUtils + +/** + * Tests for [[DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_ENABLED]]: a MERGE records + * `target_key IN (source keys)` as a read predicate so that conflict-time data skipping can + * reconcile concurrent MERGEs on disjoint keys, while concurrent changes to the same key (an + * update of the same row, or an insert of the same new key) still conflict. + * + * Each race runs the loser MERGE up to pre-commit, commits the winner MERGE, then lets the loser + * commit. The table has two files, ids [0, 50) and [50, 100), so that two MERGEs that update + * different existing rows touch different files and the scenario does not depend on same-file + * deletion-vector reconciliation. + */ +class MergeSourceKeyReadPredicateSuite extends QueryTest + with SharedSparkSession + with DeltaSQLCommandTest + with PhaseLockingTestMixin + with TransactionExecutionTestMixin { + + private val COMMITTED = "committed" + + private def tableRef(dir: File): String = s"delta.`${dir.getCanonicalPath}`" + + private def createTable(dir: File, deletionVectors: Boolean): Unit = { + val path = dir.getCanonicalPath + spark.range(0, 50).selectExpr("id", "0 AS v").repartition(1) + .write.format("delta").save(path) + spark.range(50, 100).selectExpr("id", "0 AS v").repartition(1) + .write.format("delta").mode("append").save(path) + sql(s"ALTER TABLE ${tableRef(dir)} SET TBLPROPERTIES " + + s"('delta.enableDeletionVectors' = '$deletionVectors')") + } + + private def upsert(dir: File, keys: Seq[Long], v: Int): String = { + val values = keys.map(k => s"(${k}L, $v)").mkString(", ") + s"""MERGE INTO ${tableRef(dir)} t + |USING (SELECT * FROM VALUES $values AS s(id, v)) s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET v = s.v + |WHEN NOT MATCHED THEN INSERT *""".stripMargin + } + + private def insertOnly(dir: File, keys: Seq[Long], v: Int): String = { + val values = keys.map(k => s"(${k}L, $v)").mkString(", ") + s"""MERGE INTO ${tableRef(dir)} t + |USING (SELECT * FROM VALUES $values AS s(id, v)) s + |ON t.id = s.id + |WHEN NOT MATCHED THEN INSERT *""".stripMargin + } + + private def allConfs(enabled: Boolean, maxKeys: Int = 100000): Seq[(String, String)] = Seq( + DeltaSQLConf.DELTA_CONFLICT_DETECTION_DATA_SKIPPING_ENABLED.key -> "true", + DeltaSQLConf.DELTA_CONFLICT_DETECTION_DATA_SKIPPING_VALUE_EXACT_ENABLED.key -> "true", + DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_ENABLED.key -> enabled.toString, + DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_MAX_KEYS.key -> maxKeys.toString) + + /** + * Runs `loserSql` up to pre-commit, commits `winnerSql`, then lets the loser commit. Returns + * [[COMMITTED]] or the simple class name of the concurrent-modification exception. + */ + private def race(loserSql: String, winnerSql: String): String = { + val pool = ThreadUtils.newDaemonSingleThreadExecutor(threadName = "merge-loser") + try { + val (observer, future) = runQueryWithObserver(name = "loser", pool, loserSql) + unblockUntilPreCommit(observer) + busyWaitFor(observer.phases.preparePhase.hasEntered, timeout) + sql(winnerSql) + unblockCommit(observer) + try { + ThreadUtils.awaitResult(future, Duration.Inf) + COMMITTED + } catch { + case NonFatal(e) => + Iterator.iterate[Throwable](e)(_.getCause).takeWhile(_ != null) + .collectFirst { case c: DeltaConcurrentModificationException => c } + .map(_.getClass.getSimpleName) + .getOrElse(throw e) + } + } finally { + pool.shutdownNow() + } + } + + private def ids(dir: File): Seq[Long] = + sql(s"SELECT id FROM ${tableRef(dir)}").collect().map(_.getLong(0)).toSeq + + for (dv <- Seq(true, false)) { + test(s"disjoint keys: loser commits, both MERGEs applied (deletionVectors=$dv)") { + withTempDir { dir => + createTable(dir, deletionVectors = dv) + withSQLConf(allConfs(enabled = true): _*) { + // Winner updates 10 (file [0, 50)) and inserts 1000: its new file's id range [10, 1000] + // spans the loser's keys, so only the value-exact check can prove no conflict. + assert(race(upsert(dir, Seq(60L, 500L), 2), upsert(dir, Seq(10L, 1000L), 1)) == + COMMITTED) + } + checkAnswer( + sql(s"SELECT id, v FROM ${tableRef(dir)} WHERE id IN (10, 60, 500, 1000)"), + Seq(Row(10L, 1), Row(60L, 2), Row(500L, 2), Row(1000L, 1))) + val all = ids(dir) + assert(all.size == 102 && all.distinct.size == 102, "no lost or duplicated rows") + } + } + } + + test("feature off: disjoint keys still conflict (key-only ON reads the whole table)") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + withSQLConf(allConfs(enabled = false): _*) { + assert(race(upsert(dir, Seq(60L, 500L), 2), upsert(dir, Seq(10L, 1000L), 1)) == + "ConcurrentAppendException") + } + } + } + + test("same new key inserted by both MERGEs: loser aborts, no duplicate key") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + withSQLConf(allConfs(enabled = true): _*) { + assert(race(upsert(dir, Seq(60L, 500L), 2), upsert(dir, Seq(10L, 500L), 1)) == + "ConcurrentAppendException") + } + assert(ids(dir).count(_ == 500L) == 1, "the new key must be inserted exactly once") + } + } + + test("same existing key updated by both MERGEs: loser aborts") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + withSQLConf(allConfs(enabled = true): _*) { + val outcome = race(upsert(dir, Seq(60L), 2), upsert(dir, Seq(60L), 1)) + assert(outcome != COMMITTED, "updating the same row concurrently must conflict") + } + checkAnswer(sql(s"SELECT v FROM ${tableRef(dir)} WHERE id = 60"), Row(1)) + } + } + + test("more distinct source keys than maxKeys: falls back to the existing read predicate") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + withSQLConf(allConfs(enabled = true, maxKeys = 1): _*) { + assert(race(upsert(dir, Seq(60L, 500L), 2), upsert(dir, Seq(10L, 1000L), 1)) == + "ConcurrentAppendException") + } + } + } + + test("insert-only MERGE: disjoint new keys commit, the same new key aborts") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + withSQLConf(allConfs(enabled = true): _*) { + assert(race(insertOnly(dir, Seq(500L), 2), insertOnly(dir, Seq(1000L, 10L), 1)) == + COMMITTED) + assert(race(insertOnly(dir, Seq(700L), 2), insertOnly(dir, Seq(2000L, 700L), 1)) == + "ConcurrentAppendException") + } + val all = ids(dir) + assert(all.count(_ == 500L) == 1 && all.count(_ == 1000L) == 1 && all.count(_ == 700L) == 1) + } + } + + test("NOT MATCHED BY SOURCE reads every target row: no source-key predicate, conflicts") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + val loser = + s"""MERGE INTO ${tableRef(dir)} t + |USING (SELECT * FROM VALUES (60L, 2) AS s(id, v)) s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET v = s.v + |WHEN NOT MATCHED BY SOURCE AND t.id = 1000 THEN DELETE""".stripMargin + withSQLConf(allConfs(enabled = true): _*) { + assert(race(loser, upsert(dir, Seq(10L, 1000L), 1)) != COMMITTED) + } + } + } + + // --------------------------------------------------------------------------------------------- + // Predicate construction, observed through the usage event. + // --------------------------------------------------------------------------------------------- + + private def sourceKeyEvents(f: => Unit): Seq[Map[String, Any]] = + Log4jUsageLogger.track(f) + .filter(_.tags.get("opType").contains("delta.dml.merge.sourceKeyReadPredicate")) + .map(e => JsonUtils.fromJson[Map[String, Any]](e.blob)) + + test("event: applied within maxKeys, not applied above, nulls ignored, keys de-duplicated") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + withSQLConf(allConfs(enabled = true, maxKeys = 2): _*) { + // Duplicate source keys (existing target rows, so insert-only MERGE is a no-op). + val within = sourceKeyEvents(sql(insertOnly(dir, Seq(1L, 2L, 2L), 5))) + assert(within.size == 1) + assert(within.head("applied") == true && within.head("numDistinctKeys") == 2) + + val above = sourceKeyEvents(sql(upsert(dir, Seq(1L, 2L, 3L), 6))) + assert(above.size == 1) + assert(above.head("applied") == false && above.head("numDistinctKeys") == 3) + + val nullOnly = sourceKeyEvents(sql( + s"""MERGE INTO ${tableRef(dir)} t + |USING (SELECT CAST(NULL AS BIGINT) AS id, 7 AS v) s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET v = s.v""".stripMargin)) + assert(nullOnly.size == 1 && nullOnly.head("applied") == false) + } + checkAnswer(sql(s"SELECT id, v FROM ${tableRef(dir)} WHERE id IN (1, 2, 3)"), + Seq(Row(1L, 6), Row(2L, 6), Row(3L, 6))) + } + } + + test("event: no equi-join key between target and source -> not attempted") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + withSQLConf(allConfs(enabled = true): _*) { + val events = sourceKeyEvents(sql( + s"""MERGE INTO ${tableRef(dir)} t + |USING (SELECT 60L AS id, 9 AS v) s + |ON t.id < s.id AND t.id >= 59 + |WHEN MATCHED THEN UPDATE SET v = s.v""".stripMargin)) + assert(events.isEmpty) + } + checkAnswer(sql(s"SELECT v FROM ${tableRef(dir)} WHERE id = 59"), Row(9)) + } + } + + test("multi-column key: per-column IN predicates, MERGE result unchanged") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + withSQLConf(allConfs(enabled = true): _*) { + val events = sourceKeyEvents(sql( + s"""MERGE INTO ${tableRef(dir)} t + |USING (SELECT * FROM VALUES (60L, 0, 3), (61L, 0, 3), (62L, 1, 3) AS s(id, v, nv)) s + |ON t.id = s.id AND t.v = s.v + |WHEN MATCHED THEN UPDATE SET v = s.nv""".stripMargin)) + assert(events.size == 1) + assert(events.head("applied") == true && events.head("numKeyColumns") == 2) + } + checkAnswer(sql(s"SELECT id, v FROM ${tableRef(dir)} WHERE id IN (60, 61, 62)"), + Seq(Row(60L, 3), Row(61L, 3), Row(62L, 0))) + } + } +} From 0f0a32d46dbd303d7e337cf6c96fbd1b9134fc7b Mon Sep 17 00:00:00 2001 From: Eunjin Song Date: Tue, 6 Oct 2026 22:58:54 -0700 Subject: [PATCH 2/2] [RLC] MERGE source-key read predicate: exact semi-join (DBR parity), no key cap Updated DBR study (wiki, DBR 18.2 source-size addendum): DBR never decides MERGE conflicts with a literal IN list. Its read predicate is a join against the MERGE source (DeltaTableMergeSourceReadPredicate), with no source-size cap (tested to 4M rows), and rowLevelConflictDetection.inSetThreshold is only a pre-filter. Match that: - Replace merge.sourceKeyReadPredicate.maxKeys (and its fallback) with merge.sourceKeyReadPredicate.inSetThreshold (default 10000). It only chooses the pre-filter: InSet up to the threshold, min/max range above it. - Add MergeSourceKeyMatch(targetKeys, sourceKeys), recorded with the pre-filters. It is `true` in generic paths; the value-exact and delete/read conflict tiers evaluate it as a left-semi join of the winner's rows against the (materialized or deterministic) MERGE source on the key tuple. - Exclude FLOAT/DOUBLE and non-binary-collation string keys, where a value pre-filter could disagree with the join's `=`. - Tests: 20000-key source (range pre-filter) commits on disjoint keys and aborts on a shared key; inSetThreshold 1/0 never changes the outcome; multi-column keys are matched as tuples. A mutation check (join removed) fails exactly these three tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../delta/commands/MergeIntoCommandBase.scala | 104 ++++++++++++------ .../expressions/MergeSourceKeyMatch.scala | 62 +++++++++++ .../sql/delta/sources/DeltaSQLConf.scala | 41 +++---- .../stats/ConflictDataSkippingReader.scala | 54 +++++++-- .../MergeSourceKeyReadPredicateSuite.scala | 86 ++++++++++++--- 5 files changed, 263 insertions(+), 84 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/delta/expressions/MergeSourceKeyMatch.scala diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/commands/MergeIntoCommandBase.scala b/spark/src/main/scala/org/apache/spark/sql/delta/commands/MergeIntoCommandBase.scala index 6847f218bd2..c47ad800e6b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/delta/commands/MergeIntoCommandBase.scala +++ b/spark/src/main/scala/org/apache/spark/sql/delta/commands/MergeIntoCommandBase.scala @@ -25,6 +25,7 @@ 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 @@ -39,8 +40,8 @@ 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.{AtomicType, 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 @@ -388,14 +389,21 @@ trait MergeIntoCommandBase extends LeafRunnableCommand /** * The `(target, source)` expression pairs of the `target_expr = source_expr` conjuncts of the ON - * condition, where each side is deterministic, atomic-typed and references only its own side. - * Null-safe equality is excluded: a NULL source key would have to match NULL target keys, which - * an IN predicate cannot express. + * 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) && - e.dataType.isInstanceOf[AtomicType] + 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)) { @@ -410,46 +418,70 @@ trait MergeIntoCommandBase extends LeafRunnableCommand } /** - * Read predicates `target_expr IN (distinct source keys)`, one per equi-join key of the ON - * condition, recorded so that conflict detection knows which target rows this MERGE depends on. + * 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. * - * 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). Each target row that can satisfy - * the ON condition has every equi-join key value in the set of that key's non-null source values, - * so the conjunction of the per-key IN predicates is a sound over-approximation of the rows read, - * including the rows whose absence lets a NOT MATCHED clause insert. It is also sound for MATCHED - * clause conditions, which only narrow the joined rows. + * 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. * - * Returns no predicate (the MERGE keeps its existing read predicates) when the feature is off, - * the ON condition has no usable equi-join key, a key has no non-null source value, or there are - * more than [[DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_MAX_KEYS]] distinct key tuples. The - * source is the same prepared (materialized, or deterministic) source the MERGE joins with, so - * the collected keys are the keys the join sees. + * 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] = { - val maxKeys = spark.conf.get(DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_MAX_KEYS) if (!spark.conf.get(DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_ENABLED) || - maxKeys <= 0 || notMatchedBySourceClauses.nonEmpty) { + notMatchedBySourceClauses.nonEmpty) { return Seq.empty } val keyPairs = sourceTargetEquiJoinKeys if (keyPairs.isEmpty) { return Seq.empty } - - val keyColumns = keyPairs.zipWithIndex.map { case ((_, sourceKey), i) => - Column(Alias(sourceKey, s"__merge_source_key_$i")()) - } - val keyRows = getMergeSource.df.select(keyColumns: _*).distinct().limit(maxKeys + 1).collect() - val predicates = if (keyRows.length > maxKeys) { - None - } else { - val perKey = keyPairs.zipWithIndex.map { case ((targetKey, sourceKey), i) => + 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) None else Some(InSet(targetKey, values)) + 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))) + } } - if (perKey.forall(_.isDefined)) Some(perKey.flatten) else None } recordDeltaEvent( @@ -457,10 +489,10 @@ trait MergeIntoCommandBase extends LeafRunnableCommand opType = "delta.dml.merge.sourceKeyReadPredicate", data = Map( "numKeyColumns" -> keyPairs.size, - "numDistinctKeys" -> keyRows.length, - "maxKeys" -> maxKeys, - "applied" -> predicates.isDefined)) - predicates.getOrElse(Seq.empty) + "inSetThreshold" -> inSetThreshold, + "preFilter" -> (if (useInSet) "inSet" else "range"), + "numDistinctKeys" -> (if (useInSet) keyRows.length.toLong else -1L))) + preFilters :+ MergeSourceKeyMatch(keyPairs.map(_._1), sourceKeys) } /** diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/expressions/MergeSourceKeyMatch.scala b/spark/src/main/scala/org/apache/spark/sql/delta/expressions/MergeSourceKeyMatch.scala new file mode 100644 index 00000000000..237f6e3877e --- /dev/null +++ b/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(", ")})" +} diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSQLConf.scala b/spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSQLConf.scala index 166e4f5def7..02dd550de41 100644 --- a/spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSQLConf.scala +++ b/spark/src/main/scala/org/apache/spark/sql/delta/sources/DeltaSQLConf.scala @@ -590,33 +590,34 @@ trait DeltaSQLConfBase extends DeltaSQLConfUtils { buildConf("merge.sourceKeyReadPredicate.enabled") .internal() .doc( - """When enabled, MERGE (without NOT MATCHED BY SOURCE clauses) collects the distinct - |values of the source side of each `target_expr = source_expr` equi-join conjunct of the - |ON condition and records `target_expr IN (source keys)` as an extra read predicate of - |the transaction (it is also used to prune target files for the MERGE's own scan). - |Without it, a key-only ON condition such as `t.id = s.id` registers a read predicate of - |`true`, so any concurrently added file conflicts with the MERGE. With it, conflict-time - |data skipping (conflictDetection.dataSkipping.enabled, and especially the value-exact - |tier) can prove that concurrent MERGEs on disjoint keys do not conflict, while a - |concurrent change to any of this MERGE's keys (an update, or an insert of the same new - |key) still conflicts. When the number of distinct source key tuples exceeds - |merge.sourceKeyReadPredicate.maxKeys the predicate is not recorded and MERGE behaves - |exactly as when this flag is off.""".stripMargin) + """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_MAX_KEYS = - buildConf("merge.sourceKeyReadPredicate.maxKeys") + 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 records a `target_expr IN (source keys)` read - |predicate. The keys are collected to the driver, so this bounds driver memory and the - |size of the predicate evaluated during file pruning and conflict detection. Above the - |limit the MERGE falls back to its existing read predicates. A non-positive value - |disables the feature.""".stripMargin) + |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(100000) + .createWithDefault(10000) val DELTA_PROTOCOL_DEFAULT_WRITER_VERSION = buildConf("properties.defaults.minWriterVersion") diff --git a/spark/src/main/scala/org/apache/spark/sql/delta/stats/ConflictDataSkippingReader.scala b/spark/src/main/scala/org/apache/spark/sql/delta/stats/ConflictDataSkippingReader.scala index c6ada250830..a34b2ebb0d3 100644 --- a/spark/src/main/scala/org/apache/spark/sql/delta/stats/ConflictDataSkippingReader.scala +++ b/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,11 +270,6 @@ 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 @@ -282,7 +277,7 @@ trait ConflictDataSkippingReader extends DeltaLogging { self: DataSkippingReader 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 diff --git a/spark/src/test/scala/org/apache/spark/sql/delta/MergeSourceKeyReadPredicateSuite.scala b/spark/src/test/scala/org/apache/spark/sql/delta/MergeSourceKeyReadPredicateSuite.scala index 7613a091747..97db313831d 100644 --- a/spark/src/test/scala/org/apache/spark/sql/delta/MergeSourceKeyReadPredicateSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/delta/MergeSourceKeyReadPredicateSuite.scala @@ -32,10 +32,11 @@ import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.util.ThreadUtils /** - * Tests for [[DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_ENABLED]]: a MERGE records - * `target_key IN (source keys)` as a read predicate so that conflict-time data skipping can - * reconcile concurrent MERGEs on disjoint keys, while concurrent changes to the same key (an - * update of the same row, or an insert of the same new key) still conflict. + * Tests for [[DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_ENABLED]]: a MERGE records the target + * rows whose key matches a source key as its read predicate (a pre-filter plus an exact semi-join + * against the MERGE source, like DBR) so that conflict-time data skipping can reconcile concurrent + * MERGEs on disjoint keys, however many keys there are, while concurrent changes to the same key + * (an update of the same row, or an insert of the same new key) still conflict. * * Each race runs the loser MERGE up to pre-commit, commits the winner MERGE, then lets the loser * commit. The table has two files, ids [0, 50) and [50, 100), so that two MERGEs that update @@ -79,11 +80,13 @@ class MergeSourceKeyReadPredicateSuite extends QueryTest |WHEN NOT MATCHED THEN INSERT *""".stripMargin } - private def allConfs(enabled: Boolean, maxKeys: Int = 100000): Seq[(String, String)] = Seq( + private def allConfs( + enabled: Boolean, + inSetThreshold: Int = 10000): Seq[(String, String)] = Seq( DeltaSQLConf.DELTA_CONFLICT_DETECTION_DATA_SKIPPING_ENABLED.key -> "true", DeltaSQLConf.DELTA_CONFLICT_DETECTION_DATA_SKIPPING_VALUE_EXACT_ENABLED.key -> "true", DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_ENABLED.key -> enabled.toString, - DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_MAX_KEYS.key -> maxKeys.toString) + DeltaSQLConf.MERGE_SOURCE_KEY_READ_PREDICATE_IN_SET_THRESHOLD.key -> inSetThreshold.toString) /** * Runs `loserSql` up to pre-commit, commits `winnerSql`, then lets the loser commit. Returns @@ -166,13 +169,61 @@ class MergeSourceKeyReadPredicateSuite extends QueryTest } } - test("more distinct source keys than maxKeys: falls back to the existing read predicate") { + test("more source keys than inSetThreshold: range pre-filter, the exact key match decides") { withTempDir { dir => createTable(dir, deletionVectors = true) - withSQLConf(allConfs(enabled = true, maxKeys = 1): _*) { - assert(race(upsert(dir, Seq(60L, 500L), 2), upsert(dir, Seq(10L, 1000L), 1)) == - "ConcurrentAppendException") + // 20000 even keys in [1000, 40998]: above the default threshold, so the pre-filter is the + // range [1000, 40998], which every winner key below falls into. + val loser = + s"""MERGE INTO ${tableRef(dir)} t + |USING (SELECT 1000 + 2 * id AS id, 2 AS v FROM range(20000)) s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET v = s.v + |WHEN NOT MATCHED THEN INSERT *""".stripMargin + withSQLConf(allConfs(enabled = true): _*) { + val events = sourceKeyEvents { + assert(race(loser, upsert(dir, Seq(10L, 1001L), 1)) == COMMITTED) + } + assert(events.exists(_("preFilter") == "range")) + // Control: the winner updates one of the loser's 20000 keys. + assert(race(loser.replace("2 AS v", "3 AS v"), upsert(dir, Seq(1002L), 1)) != COMMITTED) + } + val all = ids(dir) + assert(all.size == 100 + 20000 + 1 && all.distinct.size == all.size) + } + } + + test("tiny inSetThreshold only changes the pre-filter: disjoint keys still commit") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + for (threshold <- Seq(1, 0)) { + withSQLConf(allConfs(enabled = true, inSetThreshold = threshold): _*) { + // The winner's new key 300 is inside the loser's key range [60, 500 + threshold]. + val loserKeys = Seq(60L, 500L + threshold) + assert(race(upsert(dir, loserKeys, 2), upsert(dir, Seq(10L, 300L + threshold), 1)) == + COMMITTED) + } + } + assert(ids(dir).size == 100 + 4) + } + } + + test("multi-column key: tuples are matched exactly, not per column") { + withTempDir { dir => + createTable(dir, deletionVectors = true) + // The loser reads (id, v) in {(60, 0), (61, 5)}. The winner's (60, 5) passes the per-column + // pre-filters (id IN (60, 61), v IN (0, 5)) but is not one of the loser's key tuples. + val loser = + s"""MERGE INTO ${tableRef(dir)} t + |USING (SELECT * FROM VALUES (60L, 0, 7), (61L, 5, 7) AS s(id, v, nv)) s + |ON t.id = s.id AND t.v = s.v + |WHEN MATCHED THEN UPDATE SET v = s.nv + |WHEN NOT MATCHED THEN INSERT (id, v) VALUES (s.id, s.v)""".stripMargin + withSQLConf(allConfs(enabled = true): _*) { + assert(race(loser, s"INSERT INTO ${tableRef(dir)} VALUES (60, 5)") == COMMITTED) } + checkAnswer(sql(s"SELECT id, v FROM ${tableRef(dir)} WHERE id IN (60, 61)"), + Seq(Row(60L, 7), Row(60L, 5), Row(61L, 0), Row(61L, 5))) } } @@ -214,25 +265,26 @@ class MergeSourceKeyReadPredicateSuite extends QueryTest .filter(_.tags.get("opType").contains("delta.dml.merge.sourceKeyReadPredicate")) .map(e => JsonUtils.fromJson[Map[String, Any]](e.blob)) - test("event: applied within maxKeys, not applied above, nulls ignored, keys de-duplicated") { + test("event: IN pre-filter within inSetThreshold, range above, keys de-duplicated") { withTempDir { dir => createTable(dir, deletionVectors = true) - withSQLConf(allConfs(enabled = true, maxKeys = 2): _*) { + withSQLConf(allConfs(enabled = true, inSetThreshold = 2): _*) { // Duplicate source keys (existing target rows, so insert-only MERGE is a no-op). val within = sourceKeyEvents(sql(insertOnly(dir, Seq(1L, 2L, 2L), 5))) assert(within.size == 1) - assert(within.head("applied") == true && within.head("numDistinctKeys") == 2) + assert(within.head("preFilter") == "inSet" && within.head("numDistinctKeys") == 2) val above = sourceKeyEvents(sql(upsert(dir, Seq(1L, 2L, 3L), 6))) assert(above.size == 1) - assert(above.head("applied") == false && above.head("numDistinctKeys") == 3) + assert(above.head("preFilter") == "range" && above.head("numDistinctKeys") == -1) + // A NULL key matches no target row: the MERGE reads nothing and updates nothing. val nullOnly = sourceKeyEvents(sql( s"""MERGE INTO ${tableRef(dir)} t |USING (SELECT CAST(NULL AS BIGINT) AS id, 7 AS v) s |ON t.id = s.id |WHEN MATCHED THEN UPDATE SET v = s.v""".stripMargin)) - assert(nullOnly.size == 1 && nullOnly.head("applied") == false) + assert(nullOnly.size == 1) } checkAnswer(sql(s"SELECT id, v FROM ${tableRef(dir)} WHERE id IN (1, 2, 3)"), Seq(Row(1L, 6), Row(2L, 6), Row(3L, 6))) @@ -254,7 +306,7 @@ class MergeSourceKeyReadPredicateSuite extends QueryTest } } - test("multi-column key: per-column IN predicates, MERGE result unchanged") { + test("event: multi-column key, MERGE result unchanged") { withTempDir { dir => createTable(dir, deletionVectors = true) withSQLConf(allConfs(enabled = true): _*) { @@ -264,7 +316,7 @@ class MergeSourceKeyReadPredicateSuite extends QueryTest |ON t.id = s.id AND t.v = s.v |WHEN MATCHED THEN UPDATE SET v = s.nv""".stripMargin)) assert(events.size == 1) - assert(events.head("applied") == true && events.head("numKeyColumns") == 2) + assert(events.head("preFilter") == "inSet" && events.head("numKeyColumns") == 2) } checkAnswer(sql(s"SELECT id, v FROM ${tableRef(dir)} WHERE id IN (60, 61, 62)"), Seq(Row(60L, 3), Row(61L, 3), Row(62L, 0)))