You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ge...@apache.org on 2026/06/08 16:08:57 UTC
(spark) branch branch-4.2 updated: [SPARK-56995][SQL][DML] Allow dataframe caching in the DSv2 Transaction API
This is an automated email from the ASF dual-hosted git repository.
gengliangwang pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.2 by this push:
new 96401536ec6e [SPARK-56995][SQL][DML] Allow dataframe caching in the DSv2 Transaction API
96401536ec6e is described below
commit 96401536ec6eb63580420349d3d0a9b05658f94b
Author: Andreas Chatzistergiou <an...@databricks.com>
AuthorDate: Mon Jun 8 09:08:23 2026 -0700
[SPARK-56995][SQL][DML] Allow dataframe caching in the DSv2 Transaction API
Currently, the DSv2 Transaction API skips dataframe caching. This can cause significant performance regression to existing workloads. Dataframes cached prior to the transaction should be allowed to be reused within the transaction. Cache substitution during a transaction now delegates to the connector via `Transaction.registerScans`. Spark hands every materialized scan in a candidate cached subtree to the active transaction, and the connector decides whether reusing the cached snapsho [...]
Furthermore, this PR fixes an issue where the `v2TableReference` to Relation mechanism would not take into account subqueries. With the fix all pre-resolved relations in the plan (including subqueries) are un-resolved to `v2TableReference` and then resolved again at `ResolveRelations`.
### What changes were proposed in this pull request?
This transaction introduces `Transaction.registerScans` in the transaction API. Check above for more details.
### Why are the changes needed?
Without this fix cached dataframes cannot be used in transactions. As a result, the DSv2 transaction API will introduce significant performance regression for relevant workloads.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing and new tests.
### Was this patch authored or co-authored using generative AI tooling?
Opus 4.7
Closes #56121 from andreaschat-db/dsv2TransactionDFCachingFix.
Authored-by: Andreas Chatzistergiou <an...@databricks.com>
Signed-off-by: Gengliang Wang <ge...@apache.org>
(cherry picked from commit 542ea3b60926b118775c39bcf5755f2c02fa43ad)
Signed-off-by: Gengliang Wang <ge...@apache.org>
---
.../catalog/transactions/Transaction.java | 17 ++
.../spark/sql/catalyst/analysis/Analyzer.scala | 9 +
.../analysis/UnresolveRelationsInTransaction.scala | 7 +-
.../transactions/TransactionUtilsSuite.scala | 2 +
.../sql/connector/catalog/InMemoryBaseTable.scala | 22 ++
.../catalog/InMemoryRowLevelOperationTable.scala | 26 ++
.../InMemoryRowLevelOperationTableCatalog.scala | 22 +-
.../connector/catalog/InMemoryTableCatalog.scala | 8 +
.../apache/spark/sql/connector/catalog/txns.scala | 37 ++-
.../datasources/v2/DataSourceV2RelationSuite.scala | 2 +-
.../apache/spark/sql/execution/CacheManager.scala | 43 ++-
.../spark/sql/execution/QueryExecution.scala | 16 +-
.../sql/connector/MergeIntoDataFrameSuite.scala | 316 ++++++++++++++++++++-
.../sql/connector/MergeIntoTableSuiteBase.scala | 14 +-
.../sql/connector/RowLevelOperationSuiteBase.scala | 6 +-
15 files changed, 522 insertions(+), 25 deletions(-)
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/transactions/Transaction.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/transactions/Transaction.java
index 77044c6202fb..130c66d05a32 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/transactions/Transaction.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/transactions/Transaction.java
@@ -20,6 +20,7 @@ package org.apache.spark.sql.connector.catalog.transactions;
import org.apache.spark.annotation.Evolving;
import org.apache.spark.sql.connector.catalog.CatalogPlugin;
import org.apache.spark.sql.connector.catalog.TransactionalCatalogPlugin;
+import org.apache.spark.sql.connector.read.Scan;
import java.io.Closeable;
@@ -66,6 +67,22 @@ public interface Transaction extends Closeable {
*/
void abort();
+ /**
+ * Attempts to register materialized scans against this transaction's read set.
+ * <p>
+ * An example use case is cache reuse. Spark passes the scans of a candidate cached subtree
+ * for the transaction's catalog and the connector decides whether to accept them.
+ * <p>
+ * The connector must either accept all passed scans (returning {@code true} after adding
+ * the scans to the read set) or refuse (returning {@code false} without modifying
+ * the read set).
+ *
+ * @param scans the materialized scans Spark offers for registration against
+ * this transaction's read set.
+ * @return true if the connector accepts the scans; false otherwise.
+ */
+ boolean registerScans(Scan[] scans);
+
/**
* Releases any resources held by this transaction.
* <p>
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index 24d88c674c7c..651fd06c898a 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -2552,6 +2552,15 @@ class Analyzer(
case a: FunctionTableSubqueryArgumentExpression if !a.plan.resolved =>
resolveSubQuery(a, outer)(
(plan, outerAttrs) => a.copy(plan = plan, outerAttrs = outerAttrs))
+ // The subquery's plan is already resolved. Replace any V2TableReferences without
+ // re-running any analyzer rules.
+ case se: SubqueryExpression
+ if se.plan.resolved &&
+ se.plan.collectFirstWithSubqueries { case _: V2TableReference => () }.isDefined =>
+ val newPlan = se.plan.transformWithSubqueries {
+ case r: V2TableReference => relationResolution.resolveReference(r)
+ }
+ se.withNewPlan(newPlan)
}
}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnresolveRelationsInTransaction.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnresolveRelationsInTransaction.scala
index 8ee64e32376f..b0497897bbac 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnresolveRelationsInTransaction.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/UnresolveRelationsInTransaction.scala
@@ -54,7 +54,12 @@ class UnresolveRelationsInTransaction(val catalogManager: CatalogManager)
private def unresolveRelations(
plan: LogicalPlan,
catalog: CatalogPlugin): LogicalPlan = {
- plan.transform {
+ // Walk subqueries too: relations from the transaction's catalog that hide inside scalar /
+ // IN / EXISTS subqueries must also be re-resolved through the txn-aware catalog so the
+ // connector can track every read in the transaction's scope. We use `transformWithSubqueries`
+ // rather than `resolveOperators*` so that subtrees marked `analyzed = true` (e.g. from
+ // cached/pre-analyzed DataFrames or temp views) are still rewritten.
+ plan.transformWithSubqueries {
case r: DataSourceV2Relation if isLoadedFromCatalog(r, catalog) =>
V2TableReference.createForTransaction(r)
}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/transactions/TransactionUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/transactions/TransactionUtilsSuite.scala
index ee771fe3f246..c6932f416e5e 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/transactions/TransactionUtilsSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/transactions/TransactionUtilsSuite.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.catalyst.transactions
import org.apache.spark.{SparkException, SparkFunSuite}
import org.apache.spark.sql.connector.catalog.{CatalogPlugin, TransactionalCatalogPlugin}
import org.apache.spark.sql.connector.catalog.transactions.{Transaction, TransactionInfo}
+import org.apache.spark.sql.connector.read.Scan
import org.apache.spark.sql.util.CaseInsensitiveStringMap
class TransactionUtilsSuite extends SparkFunSuite {
@@ -45,6 +46,7 @@ class TransactionUtilsSuite extends SparkFunSuite {
override def commit(): Unit = { committed = true; onCommit() }
override def abort(): Unit = { aborted = true; onAbort() }
override def close(): Unit = { closed = true; onClose() }
+ override def registerScans(scans: Array[Scan]): Boolean = false
}
private def mockTransactionalCatalog(
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala
index 6b09c0090de1..8bfcfc020fa1 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala
@@ -116,6 +116,21 @@ abstract class InMemoryBaseTable(
}
}
+ // Version-aware equality: two tables refer to the same metastore entity at the same state.
+ // Fall back to reference equality when `id()` is null (no metastore identity).
+ override def equals(obj: Any): Boolean = obj match {
+ case other: InMemoryBaseTable =>
+ if (this eq other) true
+ else if (id() == null || other.id() == null) false
+ else id() == other.id() && version() == other.version()
+ case _ => false
+ }
+
+ override def hashCode(): Int = {
+ if (id() == null) System.identityHashCode(this)
+ else java.util.Objects.hash(id(), version())
+ }
+
def increaseVersion(): Unit = {
tableVersion += 1
}
@@ -520,6 +535,7 @@ abstract class InMemoryBaseTable(
if (evaluableFilters.nonEmpty) {
scan.filter(evaluableFilters)
}
+ scan.pushedFilters = _pushedFilters
recordScanEvent(_pushedFilters)
scan
}
@@ -684,6 +700,12 @@ abstract class InMemoryBaseTable(
options: CaseInsensitiveStringMap)
extends BatchScanBaseClass(_data, readSchema, tableSchema) with SupportsRuntimeFiltering {
+ // Back-pointer to the table this scan was built against.
+ val table: InMemoryBaseTable = InMemoryBaseTable.this
+
+ // The filters pushed to this scan at build time.
+ var pushedFilters: Array[Filter] = Array.empty
+
override def filterAttributes(): Array[NamedReference] = {
val scanFields = readSchema.fields.map(_.name).toSet
partitioning.flatMap(_.references)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala
index 5c0bc0b143f3..3b94f19f6fcd 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala
@@ -79,6 +79,32 @@ class InMemoryRowLevelOperationTable private (
// (operation, id, metadata, row)
var lastWriteLog: Seq[InternalRow] = Seq.empty
+ override def copy(): Table = {
+ val copied = InMemoryRowLevelOperationTable.withColumns(
+ name = name,
+ columns = columns(),
+ partitioning = partitioning,
+ properties = properties,
+ constraints = constraints,
+ tableId = id)
+ dataMap.synchronized {
+ dataMap.foreach { case (key, splits) =>
+ val copiedSplits = splits.map { bufferedRows =>
+ val copiedBufferedRows = new BufferedRows(bufferedRows.key, bufferedRows.schema)
+ copiedBufferedRows.rows ++= bufferedRows.rows.map(_.copy())
+ copiedBufferedRows
+ }
+ copied.dataMap.put(key, copiedSplits)
+ }
+ }
+ copied.commits ++= commits.map(_.copy())
+ copied.setVersionAndValidatedVersionFrom(this)
+ copied.replacedPartitions = replacedPartitions
+ copied.lastWriteInfo = lastWriteInfo
+ copied.lastWriteLog = lastWriteLog
+ copied
+ }
+
override def newRowLevelOperationBuilder(
info: RowLevelOperationInfo): RowLevelOperationBuilder = {
if (properties.getOrDefault(SUPPORTS_DELTAS, "false") == "true") {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTableCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTableCatalog.scala
index 285a2891bc93..affc4e03b404 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTableCatalog.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTableCatalog.scala
@@ -35,11 +35,29 @@ class InMemoryRowLevelOperationTableCatalog
// All transactions in order (committed and aborted), allowing per-statement
// validation in SQL scripting tests.
val observedTransactions: ArrayBuffer[Txn] = new ArrayBuffer[Txn]()
+ // Test-only knob. When true, the next transaction created by `beginTransaction` will reject
+ // register-scans calls (`registerScans` returns false unconditionally). Reset after consumed.
+ var nextTxnRejectRegisteredScansAttempt: Boolean = false
+
+ // Each `loadTable` returns a fresh snapshot pinned at the current table version (id is
+ // preserved). This is the "pin at table loading" semantics that lets version-aware
+ // `Table.equals` catch staleness: a cached relation holds a copy frozen at V1; a later load
+ // returns a copy at V2, and the two compare unequal so cache substitution fails before
+ // `Transaction.registerScans` is consulted.
+ override def loadTable(ident: Identifier): Table = {
+ liveTable(ident) match {
+ case rlot: InMemoryRowLevelOperationTable => rlot.copy()
+ case other => other
+ }
+ }
override def beginTransaction(info: TransactionInfo): Transaction = {
assert(transaction == null || transaction.currentState != Active)
- this.transaction = new Txn(new TxnTableCatalog(this))
- transaction
+ val txn = new Txn(new TxnTableCatalog(this))
+ txn.rejectRegisteredScansAttempt = nextTxnRejectRegisteredScansAttempt
+ nextTxnRejectRegisteredScansAttempt = false
+ this.transaction = txn
+ txn
}
override def createTable(ident: Identifier, tableInfo: TableInfo): Table = {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala
index d39e422a9d9e..c9a6c4acfa01 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalog.scala
@@ -71,6 +71,14 @@ class BasicInMemoryTableCatalog extends TableCatalog {
}
}
+ // Returns the underlying live instance without copying. Used by tests that need to mutate
+ // state in a way that's observable to subsequent `loadTable` callers, and by wrappers that
+ // need to propagate writes to the live state.
+ def liveTable(ident: Identifier): Table = {
+ Option(tables.get(ident)).getOrElse(
+ throw new NoSuchTableException(ident.asMultipartIdentifier))
+ }
+
// load table for writes
override def loadTable(
ident: Identifier,
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/txns.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/txns.scala
index 4b9dff5c3d78..44489a14b4fe 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/txns.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/txns.scala
@@ -24,6 +24,7 @@ import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
import org.apache.spark.sql.connector.catalog.transactions.Transaction
+import org.apache.spark.sql.connector.read.Scan
import org.apache.spark.sql.connector.write.{LogicalWriteInfo, RowLevelOperationBuilder, RowLevelOperationInfo, WriteBuilder}
import org.apache.spark.sql.sources.Filter
import org.apache.spark.sql.types.StructType
@@ -39,10 +40,39 @@ class Txn(override val catalog: TxnTableCatalog) extends Transaction {
private[this] var state: TransactionState = Active
private[this] var closed: Boolean = false
+ // Records every batch of scans the connector accepted via registerScans. Tests assert on this
+ // to confirm cache substitution went through the txn path.
+ val registeredScans: ArrayBuffer[Seq[Scan]] = ArrayBuffer.empty
+
+ // Test-only switch. When true, `registerScans` unconditionally returns false, simulating
+ // a connector that rejects all cache reuse. Used to verify Spark's cache-bypass behavior.
+ var rejectRegisteredScansAttempt: Boolean = false
+
def currentState: TransactionState = state
def isClosed: Boolean = closed
+ // Accept the batch only if every scan can be routed to a TxnTable this catalog is tracking;
+ // otherwise refuse. Staleness is handled upstream: `loadTable` returns a snapshot pinned at
+ // load, and version-aware `Table.equals` rejects cache matches whose underlying table has
+ // moved on.
+ override def registerScans(scans: Array[Scan]): Boolean = {
+ if (rejectRegisteredScansAttempt) return false
+
+ val routed = scans.toSeq.map {
+ case s: InMemoryBaseTable#InMemoryBatchScan =>
+ catalog.txnTables.values.find(_.delegate == s.table).map(_ -> s)
+ case _ => None
+ }
+ if (routed.exists(_.isEmpty)) return false
+
+ registeredScans += scans.toSeq
+ routed.flatten.foreach { case (txnTable, s) =>
+ txnTable.scanEvents += s.pushedFilters
+ }
+ true
+ }
+
override def commit(): Unit = {
if (closed) throw new IllegalStateException("Can't commit, already closed")
if (state == Aborted) throw new IllegalStateException("Can't commit, already aborted")
@@ -86,6 +116,9 @@ class TxnTable(
// don't false-positive on the TxnTable wrapper having a different UUID.
override val id: String = delegate.id
+ // The starting version should be the delegate version.
+ setVersion(delegate.version())
+
// Preserve column IDs from the delegate so that column ID validation can correctly detect
// drop-and-re-add scenarios (different IDs) and pass when columns are unchanged (same IDs).
// Uses assignMissingIds to keep the delegate's IDs for existing columns while assigning
@@ -157,7 +190,9 @@ class TxnTableCatalog(delegate: InMemoryRowLevelOperationTableCatalog) extends T
// table.
override def loadTable(ident: Identifier): Table = {
tables.computeIfAbsent(ident, _ => {
- val table = delegate.loadTable(ident).asInstanceOf[InMemoryRowLevelOperationTable]
+ // Wrap the live underlying instance (not a snapshot copy from loadTable) so commits
+ // propagate back to the catalog's authoritative state.
+ val table = delegate.liveTable(ident).asInstanceOf[InMemoryRowLevelOperationTable]
new TxnTable(table, table.schema(), this)
})
}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala
index 10ec9efca4ab..a1f1c85fe4d5 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2RelationSuite.scala
@@ -43,7 +43,7 @@ class DataSourceV2RelationSuite extends SparkFunSuite {
rowCount = Some(10),
colStats = Map(
"id" -> idColStat,
- // "extra" is not in schema — should be silently skipped
+ // "extra" is not in schema, should be silently skipped
"extra" -> CatalogColumnStat(distinctCount = Some(5))))
val v2Stats = DataSourceV2Relation.v1StatsToV2Stats(catalogStats, schema)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala
index 3f92f24156d3..345c1d5d635f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/CacheManager.scala
@@ -35,11 +35,12 @@ import org.apache.spark.sql.classic.{Dataset, SparkSession}
import org.apache.spark.sql.connector.catalog.CatalogPlugin
import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.{IdentifierHelper, MultipartIdentifierHelper}
import org.apache.spark.sql.connector.catalog.Identifier
+import org.apache.spark.sql.connector.catalog.transactions.Transaction
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.execution.columnar.InMemoryRelation
import org.apache.spark.sql.execution.command.CommandUtils
import org.apache.spark.sql.execution.datasources.{FileIndex, HadoopFsRelation, LogicalRelation, LogicalRelationWithTable}
-import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, ExtractV2CatalogAndIdentifier, ExtractV2Table, FileTable, V2TableRefreshUtil}
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2Relation, ExtractV2CatalogAndIdentifier, ExtractV2Table, FileTable, V2TableRefreshUtil}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.storage.StorageLevel
import org.apache.spark.storage.StorageLevel.MEMORY_AND_DISK
@@ -479,8 +480,10 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper {
lookupCachedDataInternal(normalized)
}
- private def lookupCachedDataInternal(plan: LogicalPlan): Option[CachedData] = {
- val result = cachedData.find(cd => plan.sameResult(cd.plan))
+ private def lookupCachedDataInternal(
+ plan: LogicalPlan,
+ canUse: CachedData => Boolean = _ => true): Option[CachedData] = {
+ val result = cachedData.find(cd => plan.sameResult(cd.plan) && canUse(cd))
if (result.isDefined) {
CacheManager.logCacheOperation(log"Dataframe cache hit for input plan:" +
log"\n${MDC(QUERY_PLAN, plan)} matched with cache entry:" +
@@ -489,16 +492,40 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper {
result
}
+ // Decides whether the cached entry can be substituted into a plan being executed inside
+ // the given transaction. Collects only the scans whose table belongs to the transaction's
+ // catalog and asks the connector whether reusing the cached snapshot is compatible with its
+ // isolation contract. Note, this function has a side effect of mutating the read set of the
+ // connector.
+ private def validateCachedEntryForTransaction(cd: CachedData, txn: Transaction): Boolean = {
+ val txnCatalogName = txn.catalog().name()
+ val txnTables = cd.cachedRepresentation.cacheBuilder.logicalPlan.collectWithSubqueries {
+ case r: DataSourceV2Relation if r.catalog.exists(_.name() == txnCatalogName) => r.table
+ }.toSet
+ if (txnTables.isEmpty) return true
+
+ val scans = collectWithSubqueries(cd.cachedRepresentation.cacheBuilder.cachedPlan) {
+ case b: BatchScanExec if txnTables.contains(b.table) => b.scan
+ }
+ scans.nonEmpty && txn.registerScans(scans.toArray)
+ }
+
/**
* Replaces segments of the given logical plan with cached versions where possible. The input
* plan must be normalized.
+ *
+ * @param plan the plan to rewrite.
+ * @param canUse predicate filtering which cached entries are eligible for substitution.
+ * Defaults to accepting any entry.
*/
- private[sql] def useCachedData(plan: LogicalPlan): LogicalPlan = {
+ private[sql] def useCachedData(
+ plan: LogicalPlan,
+ canUse: CachedData => Boolean = _ => true): LogicalPlan = {
val newPlan = plan transformDown {
case command: Command => command
case currentFragment =>
- lookupCachedDataInternal(currentFragment).map { cached =>
+ lookupCachedDataInternal(currentFragment, canUse).map { cached =>
// After cache lookup, we should still keep the hints from the input plan.
val hints = EliminateResolvedHint.extractHintsFromPlan(currentFragment)._2
val cachedPlan = cached.cachedRepresentation.withOutput(currentFragment.output)
@@ -511,7 +538,7 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper {
}
val result = newPlan.transformAllExpressionsWithPruning(_.containsPattern(PLAN_EXPRESSION)) {
- case s: SubqueryExpression => s.withNewPlan(useCachedData(s.plan))
+ case s: SubqueryExpression => s.withNewPlan(useCachedData(s.plan, canUse))
}
if (result.fastEquals(plan)) {
@@ -527,6 +554,10 @@ class CacheManager extends Logging with AdaptiveSparkPlanHelper {
result
}
+ /** Transaction-aware variant of [[useCachedData]]. */
+ private[sql] def useCachedData(plan: LogicalPlan, txn: Transaction): LogicalPlan =
+ useCachedData(plan, validateCachedEntryForTransaction(_, txn))
+
/**
* Tries to re-cache all the cache entries that contain `resourcePath` in one or more
* `HadoopFsRelation` node(s) as part of its logical plan.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
index 3ca3032be485..d843af844d94 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
@@ -291,15 +291,13 @@ class QueryExecution(
assertAnalyzed()
assertSupported()
- // During a transaction, skip cache substitution. This is to avoid replacing relations
- // loaded by the transactional catalog with potentially stale relations cached before
- // the transaction was active.
- if (transactionOpt.isDefined) {
- normalized
- } else {
- // Clone the plan to avoid sharing the plan instance between different stages like
- // analyzing, optimizing and planning.
- sparkSession.sharedState.cacheManager.useCachedData(normalized.clone())
+ // Clone the plan to avoid sharing the plan instance between different stages like
+ // analyzing, optimizing and planning.
+ val planToRewrite = normalized.clone()
+ val cacheManager = sparkSession.sharedState.cacheManager
+ transactionOpt match {
+ case Some(txn) => cacheManager.useCachedData(planToRewrite, txn)
+ case None => cacheManager.useCachedData(planToRewrite)
}
}
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoDataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoDataFrameSuite.scala
index d2dfae13b511..cb59ce80328d 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoDataFrameSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoDataFrameSuite.scala
@@ -20,7 +20,7 @@ package org.apache.spark.sql.connector
import org.apache.spark.sql.{sources, Column, Row}
import org.apache.spark.sql.AnalysisException
import org.apache.spark.sql.classic.MergeIntoWriter
-import org.apache.spark.sql.connector.catalog.{Aborted, Committed}
+import org.apache.spark.sql.connector.catalog.{Aborted, Committed, InMemoryBaseTable, InMemoryRowLevelOperationTableCatalog}
import org.apache.spark.sql.connector.catalog.Identifier
import org.apache.spark.sql.connector.catalog.TableInfo
import org.apache.spark.sql.functions._
@@ -1233,4 +1233,318 @@ class MergeIntoDataFrameSuite extends RowLevelOperationSuiteBase {
assert(e.message.contains("incompatible changes to table `cat`.`ns1`.`source_table`"))
}
}
+
+ // Cache-substitution tests for the txn path: connector approves stale-free cached scans via
+ // Transaction.registerScans.
+ test("cached merge source is reused when the table is unchanged since caching") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ val tableVersionBeforeCache = table.version()
+
+ val sourceDF = spark.table(tableNameAsString).where("salary < 250").as("source")
+ sourceDF.cache()
+ sourceDF.count()
+
+ assert(table.version() == tableVersionBeforeCache, "sanity: caching does not bump version")
+
+ val (txn, txnTables) = executeTransaction {
+ sourceDF
+ .mergeInto(tableNameAsString, $"source.pk" === targetTableCol("pk"))
+ .whenMatched()
+ .update(Map("salary" -> targetTableCol("salary").plus(1)))
+ .merge()
+ }
+
+ assert(txn.currentState == Committed)
+ assert(txn.isClosed)
+ assert(txnTables.size == 1)
+ assert(table.version() == "2")
+ assert(txn.registeredScans.nonEmpty, "registerScans should have accepted the cached scan")
+
+ val targetTxnTable = txnTables(tableNameAsString)
+ assert(targetTxnTable.scanEvents.size == 4)
+ val sourceFilterScans = targetTxnTable.scanEvents.flatten.count {
+ case sources.LessThan("salary", 250) => true
+ case _ => false
+ }
+ assert(sourceFilterScans == 2,
+ s"expected two salary < 250 scan events, got " +
+ s"${targetTxnTable.scanEvents.map(_.toSeq).mkString("[", ", ", "]")}")
+ val targetScans = targetTxnTable.scanEvents.count(_.isEmpty)
+ assert(targetScans == 2,
+ s"expected two target-side scans with no pushed filters, got " +
+ s"${targetTxnTable.scanEvents.map(_.toSeq).mkString("[", ", ", "]")}")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Seq(
+ Row(1, 101, "hr"),
+ Row(2, 201, "software"),
+ Row(3, 300, "hr")))
+ }
+
+ test("cached merge source is dropped when the table version moves on between caching and MERGE") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ val sourceDF = spark.table(tableNameAsString).where("salary < 250").as("source")
+ sourceDF.cache()
+ sourceDF.count()
+ val versionAtCache = table.version()
+
+ // Bump the version directly to simulate an out-of-band committer. A Spark-side write would
+ // also bump the version but would trigger CacheManager.refreshCache and defeat the test.
+ table.increaseVersion()
+ assert(table.version() != versionAtCache, "sanity: bump should change the version")
+
+ val (txn, txnTables) = executeTransaction {
+ sourceDF
+ .mergeInto(tableNameAsString, $"source.pk" === targetTableCol("pk"))
+ .whenMatched()
+ .update(Map("salary" -> targetTableCol("salary").plus(1)))
+ .merge()
+ }
+
+ assert(txn.currentState == Committed)
+ assert(txn.isClosed)
+ assert(txnTables.size == 1)
+ assert(table.version() == "3")
+ assert(txn.registeredScans.isEmpty, "registerScans should refuse the stale cached scan")
+
+ val targetTxnTable = txnTables(tableNameAsString)
+ assert(targetTxnTable.scanEvents.size == 4)
+ val sourceFilterScans = targetTxnTable.scanEvents.flatten.count {
+ case sources.LessThan("salary", 250) => true
+ case _ => false
+ }
+ assert(sourceFilterScans == 2,
+ s"expected two salary<250 scan events (fresh, since cache was bypassed), got " +
+ s"${targetTxnTable.scanEvents.map(_.toSeq).mkString("[", ", ", "]")}")
+ val targetScans = targetTxnTable.scanEvents.count(_.isEmpty)
+ assert(targetScans == 2,
+ s"expected two target-side scans with no pushed filters, got " +
+ s"${targetTxnTable.scanEvents.map(_.toSeq).mkString("[", ", ", "]")}")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Seq(
+ Row(1, 101, "hr"),
+ Row(2, 201, "software"),
+ Row(3, 300, "hr")))
+ }
+
+ test("cached source from outside the txn catalog is reused without consulting registerScans") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ // Source has no DataSourceV2Relation, so validateCachedEntryForTransaction's txnTables set
+ // is empty and the cache is accepted without consulting registerScans.
+ val sourceDF = spark.range(2)
+ .select(
+ (col("id") + 1).cast(IntegerType).as("pk"),
+ lit(999).cast(IntegerType).as("salary"),
+ lit("hr").cast(StringType).as("dep"))
+ .as("source")
+ sourceDF.cache()
+ sourceDF.count()
+
+ val (txn, txnTables) = executeTransaction {
+ sourceDF
+ .mergeInto(tableNameAsString, $"source.pk" === targetTableCol("pk"))
+ .whenMatched()
+ .update(Map("salary" -> $"source.salary"))
+ .merge()
+ }
+
+ assert(txn.currentState == Committed)
+ assert(txn.isClosed)
+ assert(txnTables.size == 1)
+ assert(table.version() == "2")
+ assert(txn.registeredScans.isEmpty,
+ "registerScans should not be consulted when the cached subtree has no txn-catalog reads")
+
+ val targetTxnTable = txnTables(tableNameAsString)
+ val targetScans = targetTxnTable.scanEvents.count(_.isEmpty)
+ assert(targetScans == 2,
+ s"expected two target-side scans, got " +
+ s"${targetTxnTable.scanEvents.map(_.toSeq).mkString("[", ", ", "]")}")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Seq(
+ Row(1, 999, "hr"),
+ Row(2, 999, "software")))
+ }
+
+ test("cached relation inside an IN-subquery is substituted via useCachedData recursion") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ // The cached entry sits inside an IN-subquery, so substitution requires useCachedData's
+ // recursion into SubqueryExpression - transformDown alone wouldn't reach it.
+ val lookupName = "cat.ns1.cache_lookup"
+ withTable(lookupName) {
+ sql(s"CREATE TABLE $lookupName (pk INT) USING foo")
+ sql(s"INSERT INTO $lookupName VALUES (1), (2)")
+
+ val lookupDF = spark.table(lookupName)
+ lookupDF.cache()
+ lookupDF.count()
+ lookupDF.createOrReplaceTempView("cache_lookup_view")
+
+ val sourceDF = spark.table(tableNameAsString)
+ .where("pk IN (SELECT pk FROM cache_lookup_view)")
+ .as("source")
+
+ val (txn, txnTables) = executeTransaction {
+ sourceDF
+ .mergeInto(tableNameAsString, $"source.pk" === targetTableCol("pk"))
+ .whenMatched()
+ .update(Map("salary" -> targetTableCol("salary").plus(1)))
+ .merge()
+ }
+
+ assert(txn.currentState == Committed)
+ assert(txn.isClosed)
+ // The helper `txnTables` is built from physical BatchScanExec nodes in the executed plan;
+ // since cache_lookup is cache-served, it has no BatchScanExec and doesn't appear here.
+ assert(txnTables.size == 1)
+ // Both the target (test_table) and the IN-subquery's lookup (cache_lookup) must be
+ // tracked through the txn catalog via TxnTableCatalog.loadTable. This is the diagnostic
+ // that surfaces UnresolveRelationsInTransaction subquery-walking gaps.
+ assert(txn.catalog.txnTables.size == 2)
+ assert(table.version() == "2")
+ // Every registered scan must point at the lookup (the only cached subtree in this plan).
+ val lookupIdent = Identifier.of(Array("ns1"), "cache_lookup")
+ val lookupTable = catalog.loadTable(lookupIdent)
+ assert(txn.registeredScans.flatten.collect {
+ case s: InMemoryBaseTable#InMemoryBatchScan => s.table
+ }.distinct == Seq(lookupTable))
+ val lookupTxnTable = txn.catalog.txnTables(lookupIdent)
+ assert(lookupTxnTable.scanEvents.nonEmpty,
+ "lookup scan events should be recorded via registerScans")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Seq(
+ Row(1, 101, "hr"),
+ Row(2, 201, "software"),
+ Row(3, 300, "hr")))
+ }
+ }
+
+ test("connector rejecting registerScans causes cache bypass") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ val sourceDF = spark.table(tableNameAsString).where("salary < 250").as("source")
+ sourceDF.cache()
+ sourceDF.count()
+
+ // Force the connector to reject `registerScans` (return false) so Spark must bypass
+ // the cache.
+ catalog.nextTxnRejectRegisteredScansAttempt = true
+
+ val (txn, txnTables) = executeTransaction {
+ sourceDF
+ .mergeInto(tableNameAsString, $"source.pk" === targetTableCol("pk"))
+ .whenMatched()
+ .update(Map("salary" -> targetTableCol("salary").plus(1)))
+ .merge()
+ }
+
+ assert(txn.currentState == Committed)
+ assert(txn.isClosed)
+ assert(txnTables.size == 1)
+ assert(table.version() == "2")
+ assert(txn.registeredScans.isEmpty, "rejected registerScans attempt")
+
+ val targetTxnTable = txnTables(tableNameAsString)
+ assert(targetTxnTable.scanEvents.size == 4)
+ val sourceFilterScans = targetTxnTable.scanEvents.flatten.count {
+ case sources.LessThan("salary", 250) => true
+ case _ => false
+ }
+ assert(sourceFilterScans == 2,
+ s"expected two salary < 250 scan events (fresh, since cache was bypassed), got " +
+ s"${targetTxnTable.scanEvents.map(_.toSeq).mkString("[", ", ", "]")}")
+ val targetScans = targetTxnTable.scanEvents.count(_.isEmpty)
+ assert(targetScans == 2,
+ s"expected two target-side scans, got " +
+ s"${targetTxnTable.scanEvents.map(_.toSeq).mkString("[", ", ", "]")}")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Seq(
+ Row(1, 101, "hr"),
+ Row(2, 201, "software"),
+ Row(3, 300, "hr")))
+ }
+
+ test("registerScans receives only scans for the transaction's catalog") {
+ withSQLConf("spark.sql.catalog.cat2" ->
+ classOf[InMemoryRowLevelOperationTableCatalog].getName) {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ withTable("cat2.ns.lookup") {
+ sql("CREATE TABLE cat2.ns.lookup (pk INT) USING foo")
+ sql("INSERT INTO cat2.ns.lookup VALUES (1), (2)")
+
+ // Cache a subtree that joins a `cat` table (the txn catalog) with a `cat2` table.
+ val sourceDF = spark.table(tableNameAsString)
+ .join(spark.table("cat2.ns.lookup"), "pk")
+ .as("source")
+ sourceDF.cache()
+ sourceDF.count()
+
+ val (txn, txnTables) = executeTransaction {
+ sourceDF
+ .mergeInto(tableNameAsString, $"source.pk" === targetTableCol("pk"))
+ .whenMatched()
+ .update(Map("salary" -> targetTableCol("salary").plus(1)))
+ .merge()
+ }
+
+ assert(txn.currentState == Committed)
+ assert(txn.isClosed)
+ // Only the cat-side target is tracked in the txn's catalog; cat2.lookup goes through
+ // a different (non-transactional) catalog and must not appear here.
+ assert(txnTables.size == 1)
+ assert(txn.catalog.txnTables.size == 1)
+ assert(table.version() == "2")
+ val passed = txn.registeredScans.flatten.collect {
+ case s: InMemoryBaseTable#InMemoryBatchScan => s
+ }
+ assert(passed.nonEmpty, "registerScans should have been consulted")
+ // The filter should pass only the cat-side scan; the cat2 scan must be filtered out.
+ val catTable = catalog.loadTable(ident)
+ assert(passed.forall(_.table.id() == catTable.id()),
+ "only the txn-catalog's scan should reach registerScans")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Seq(
+ Row(1, 101, "hr"),
+ Row(2, 201, "software")))
+ }
+ }
+ }
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala
index f37a614f99b5..126b84b507ca 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/MergeIntoTableSuiteBase.scala
@@ -1198,10 +1198,18 @@ abstract class MergeIntoTableSuiteBase extends RowLevelOperationSuiteBase
assert(txn.currentState == Committed)
assert(txn.isClosed)
- // both target and source must have been read through the transaction catalog
- assert(txnTables.size == 2)
+ // The source TxnTable is still created during analysis (the txn catalog routes the
+ // load), but cache substitution reuses the cached scan instead of issuing a fresh
+ // BatchScanExec for the source, so only the target appears in the executed plan.
+ assert(txn.catalog.txnTables.size == 2)
+ assert(txnTables.size == 1)
+ assert(txnTables.contains(tableNameAsString))
assert(table.version() == "2")
- assert(txnTables(sourceNameAsString).scanEvents.nonEmpty)
+ // The connector accepted the cached source scan via registerScans, which also
+ // records the scan as a read event against the source's TxnTable.
+ assert(txn.registeredScans.nonEmpty)
+ val sourceTxnTable = txn.catalog.txnTables.values.find(_.name == sourceNameAsString).get
+ assert(sourceTxnTable.scanEvents.nonEmpty)
assert(txnTables(tableNameAsString).scanEvents.nonEmpty)
checkAnswer(
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala
index 379e7ba755d9..0c465969e347 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala
@@ -51,6 +51,7 @@ abstract class RowLevelOperationSuiteBase
}
after {
+ catalog.nextTxnRejectRegisteredScansAttempt = false
spark.sessionState.catalogManager.reset()
spark.sessionState.conf.unsetConf("spark.sql.catalog.cat")
}
@@ -94,7 +95,10 @@ abstract class RowLevelOperationSuiteBase
}
protected def table: InMemoryRowLevelOperationTable = {
- catalog.loadTable(ident).asInstanceOf[InMemoryRowLevelOperationTable]
+ // Use liveTable, not loadTable, because loadTable returns a snapshot copy.
+ // Tests mutate state through this accessor (e.g., `table.increaseVersion()`) and need
+ // those mutations to be observable to subsequent loads.
+ catalog.liveTable(ident).asInstanceOf[InMemoryRowLevelOperationTable]
}
protected def createTable(schemaString: String): Unit = {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org