You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/01 22:29:53 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6421-c593a42048ccfb8773a202a3928e14030ca74b0f created (now 2d515a8d6b)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6421-c593a42048ccfb8773a202a3928e14030ca74b0f
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 2d515a8d6b fix: keep Spark's cache scan for a relation whose cached plan records observed metrics (#6421)
This branch includes the following new commits:
new 2d515a8d6b fix: keep Spark's cache scan for a relation whose cached plan records observed metrics (#6421)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: fix: keep Spark's cache scan for a relation whose cached plan records observed metrics (#6421)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6421-c593a42048ccfb8773a202a3928e14030ca74b0f
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 2d515a8d6b8073f81dc40b8817b235cea17a4f41
Author: Andy Grove <ag...@apache.org>
AuthorDate: Thu Oct 1 22:29:31 2026 +0000
fix: keep Spark's cache scan for a relation whose cached plan records observed metrics (#6421)
* fix: keep Spark's cache scan for a relation whose cached plan records observed metrics
Spark collects Dataset.observe metrics after a query with
CollectMetricsExec.collect, which reaches the metrics recorded inside a cached
plan only through an InMemoryTableScanExec over it. With Comet's native cache
scan in its place those metrics were lost: QueryExecution.observedMetrics came
back empty, Observation.get returned an empty map on Spark 3.5 and later, and
on Spark 3.4 it never returned.
Keep Spark's scan for such a relation, with a fallback reason. The data stays
in Comet's format and is read through the existing fallback path. The check
walks the cached plan the way CollectMetricsExec.collect does, through
subqueries, adaptive plans and nested caches.
* test: cover a cache over caches that record metrics, and name the fallback condition
---
docs/source/user-guide/latest/in-memory-cache.md | 4 ++
.../org/apache/comet/rules/CometExecRule.scala | 15 +++++-
.../sql/comet/CometInMemoryTableScanExec.scala | 25 ++++++++-
.../comet/exec/CometInMemoryCacheSuite.scala | 63 ++++++++++++++++++++--
4 files changed, 100 insertions(+), 7 deletions(-)
diff --git a/docs/source/user-guide/latest/in-memory-cache.md b/docs/source/user-guide/latest/in-memory-cache.md
index 8d2ccd2521..8de4759957 100644
--- a/docs/source/user-guide/latest/in-memory-cache.md
+++ b/docs/source/user-guide/latest/in-memory-cache.md
@@ -61,6 +61,10 @@ under one setting stays readable after the setting changes. Turning
`spark.comet.exec.inMemoryCache.enabled` off at runtime only sends cached scans back to Spark's
execution path; the cached data stays readable either way.
+A relation whose cached plan records observed metrics, from `Dataset.observe`, is still stored in
+Comet's format but is scanned by Spark's `InMemoryTableScanExec`, because Spark collects those
+metrics only through that scan.
+
## Storage format
Each cached batch is stored as a single Arrow IPC record batch message and its body.
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index aea2798af4..c1aa237970 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -386,8 +386,13 @@ case class CometExecRule(session: SparkSession)
val cometCacheFormat = usesCometCacheSerializer &&
ArrowCachedBatchSerializer.supportsSchema(scan.relation.output)
val nativeCacheEnabled = CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.get(conf)
+ // Walks the cached plan, so it is lazy: only consulted once the native scan is otherwise
+ // possible. See CometInMemoryTableScanExec.recordsObservedMetrics.
+ lazy val recordsObservedMetrics =
+ CometInMemoryTableScanExec.recordsObservedMetrics(scan.relation)
+ val nativeScan = nativeCacheEnabled && cometCacheFormat && !recordsObservedMetrics
- if (nativeCacheEnabled && cometCacheFormat) {
+ if (nativeScan) {
convertToComet(scan, CometInMemoryTableScanExec).getOrElse(scan)
} else {
// The native cache scan is not available for this relation. Record why, then take the
@@ -398,7 +403,7 @@ case class CometExecRule(session: SparkSession)
scan,
s"Comet in-memory cache requires ${classOf[ArrowCachedBatchSerializer].getName} " +
s"but this relation was cached with ${serializer.getClass.getName}")
- } else if (nativeCacheEnabled) {
+ } else if (nativeCacheEnabled && !cometCacheFormat) {
val unsupported = scan.relation.output
.filterNot(a => ArrowCachedBatchSerializer.supportsType(a.dataType))
.map(a => s"${a.name}: ${a.dataType.simpleString}")
@@ -406,6 +411,12 @@ case class CometExecRule(session: SparkSession)
scan,
"Comet in-memory cache does not support the type of these cached columns, so the " +
s"relation was cached in Spark's default format: ${unsupported.mkString(", ")}")
+ } else if (nativeCacheEnabled && recordsObservedMetrics) {
+ withFallbackReason(
+ scan,
+ "Comet in-memory cache does not scan a relation whose cached plan records " +
+ "Dataset.observe metrics, because Spark collects those metrics only through " +
+ "InMemoryTableScanExec")
} else if (usesCometCacheSerializer) {
withFallbackReason(
scan,
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala
index 7af36cac66..d3f1ac03d3 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometInMemoryTableScanExec.scala
@@ -26,8 +26,9 @@ import org.apache.spark.sql.catalyst.expressions.Attribute
import org.apache.spark.sql.catalyst.plans.logical.Statistics
import org.apache.spark.sql.columnar.{CachedBatch, CachedBatchSerializer}
import org.apache.spark.sql.comet.shims.ShimCometInMemoryTableScanExec
-import org.apache.spark.sql.execution.SparkPlan
-import org.apache.spark.sql.execution.columnar.{CachedRDDBuilder, InMemoryTableScanExec}
+import org.apache.spark.sql.execution.{CollectMetricsExec, SparkPlan}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.columnar.{CachedRDDBuilder, InMemoryRelation, InMemoryTableScanExec}
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
import org.apache.spark.sql.vectorized.ColumnarBatch
@@ -170,4 +171,24 @@ object CometInMemoryTableScanExec extends CometOperatorSerde[InMemoryTableScanEx
op.output))
}
+ /**
+ * Whether `relation`'s cached plan records observed metrics, from `Dataset.observe`.
+ *
+ * Spark collects those metrics once a query finishes, with `CollectMetricsExec.collect`, and
+ * that reaches the ones recorded inside a cached plan only through an `InMemoryTableScanExec`
+ * over it. It does not know this node, so replacing the scan of such a relation leaves its
+ * metrics empty, and on Spark 3.4 leaves `Observation.get` waiting for good. The walk mirrors
+ * `CollectMetricsExec.collect`, through subqueries, adaptive plans and nested caches.
+ */
+ def recordsObservedMetrics(relation: InMemoryRelation): Boolean =
+ ObservedMetrics.recordedIn(relation.cachedPlan)
+
+ private object ObservedMetrics extends AdaptiveSparkPlanHelper {
+ def recordedIn(plan: SparkPlan): Boolean =
+ collectWithSubqueries(plan) {
+ case _: CollectMetricsExec => true
+ case scan: InMemoryTableScanExec => recordedIn(scan.relation.cachedPlan)
+ }.contains(true)
+ }
+
}
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
index c963de3589..ecce906752 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -22,6 +22,9 @@ package org.apache.comet.exec
import java.{util => ju}
import java.nio.charset.StandardCharsets
+import scala.concurrent.{Await, Future}
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.duration.DurationInt
import scala.jdk.CollectionConverters._
import org.apache.arrow.compression.ZstdCompressionCodec
@@ -31,7 +34,7 @@ import org.apache.arrow.vector.compression.{CompressionCodec, CompressionUtil, N
import org.apache.arrow.vector.types.pojo.ArrowType
import org.apache.spark.CometDriverPlugin
import org.apache.spark.SparkConf
-import org.apache.spark.sql.{CometTestBase, QueryTest, Row}
+import org.apache.spark.sql.{CometTestBase, Observation, QueryTest, Row}
import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeReference, EqualTo, Expression, GreaterThan, GreaterThanOrEqual, LessThan, Literal}
import org.apache.spark.sql.columnar.{CachedBatch, SimpleMetricsCachedBatch}
import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometInMemoryTableScanExec, CometSortExec, CometSortMergeJoinExec}
@@ -42,13 +45,13 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AQEShuffl
import org.apache.spark.sql.execution.columnar.{CometInMemoryRelationHelper, InMemoryRelation, InMemoryTableScanExec}
import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec}
-import org.apache.spark.sql.functions.max
+import org.apache.spark.sql.functions.{count, lit, max, min, sum}
import org.apache.spark.sql.internal.{SQLConf, StaticSQLConf}
import org.apache.spark.sql.types._
import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
import org.apache.spark.storage.StorageLevel
-import org.apache.comet.{CometArrowAllocator, CometConf}
+import org.apache.comet.{CometArrowAllocator, CometConf, ExtendedExplainInfo}
import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark40Plus}
import org.apache.comet.vector.{CometPlainVector, CometVector}
@@ -2614,4 +2617,58 @@ class CometInMemoryCacheSuite extends CometTestBase {
spark.catalog.clearCache()
}
}
+
+ test("Comet in-memory cache keeps the observed metrics recorded in a cached plan") {
+ // Spark collects the metrics of an observe() inside a cached plan only through an
+ // InMemoryTableScanExec over it, so the scan of such a relation has to stay Spark's. Replaced,
+ // the metrics come back empty, and on Spark 3.4 Observation.get never returns. Nested the way
+ // SPARK-35695's test nests it, with a shuffle in the inner cached plan so that AQE plans it,
+ // under one more cache that records no metrics of its own. The metrics read through that top
+ // scan are only found by following it into the caches it reads.
+ withAQECache {
+ val df = spark
+ .range(0, 100, 1, 2)
+ .repartition(4)
+ .observe("inner_event", count(lit(1)).as("rows"), max($"id").as("max_id"))
+ .persist()
+ .observe("outer_event", min($"id").as("min_id"))
+ .persist()
+ .filter($"id" > 10)
+ .persist()
+ df.collect()
+ assert(
+ df.queryExecution.observedMetrics ==
+ Map("inner_event" -> Row(100L, 99L), "outer_event" -> Row(0L)))
+ val plan = df.queryExecution.executedPlan
+ assert(collect(plan) { case s: CometInMemoryTableScanExec => s }.isEmpty)
+ assert(
+ new ExtendedExplainInfo()
+ .generateExtendedInfo(plan)
+ .contains("records Dataset.observe metrics"))
+ // Still stored in Comet's format: only the scan changes.
+ assert(
+ spark.sharedState.cacheManager
+ .lookupCachedData(df)
+ .get
+ .cachedRepresentation
+ .cacheBuilder
+ .cachedColumnBuffers
+ .map(_.getClass.getName)
+ .distinct()
+ .collect()
+ .sameElements(Array("org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")))
+
+ val observation = Observation("cached_observation")
+ val observed = spark.range(10).observe(observation, sum($"id").as("total")).persist()
+ observed.collect()
+ // Bounded, so that a regression fails here rather than hanging the suite on Spark 3.4.
+ assert(Await.result(Future(observation.get), 1.minute) == Map("total" -> 45L))
+
+ val plain = spark.range(0, 100, 1, 2).persist()
+ plain.collect()
+ assert(collect(plain.queryExecution.executedPlan) { case s: CometInMemoryTableScanExec =>
+ s
+ }.nonEmpty)
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org