feat: enable Comet's in-memory cache by default

docs/source/user-guide/latest/in-memory-cache.md

@@ -24,13 +24,13 @@ format that Comet operators read directly. Without it, a cached table is stored
 format and every scan of it has to convert each batch before Comet can continue, which shows up in
 the plan as a `CometSparkColumnarToColumnar` above the cache scan.
 
-This feature is **experimental and disabled by default**. Turn it on at startup, alongside the rest
-of Comet's configuration:
+This feature is **experimental and enabled by default**. To turn it off, set the config at startup,
+alongside the rest of Comet's configuration:
 
 ```shell
 $SPARK_HOME/bin/spark-shell \
     ... \
-    --conf spark.comet.exec.inMemoryCache.enabled=true
+    --conf spark.comet.exec.inMemoryCache.enabled=false
 ```
 
 It has to be set before the `SparkContext` starts. Comet's driver plugin chooses
@@ -48,6 +48,9 @@ With Comet's serializer installed as `spark.sql.cache.serializer`:
 - Cached tables are scanned by `CometInMemoryTableScan`, which feeds Comet operators directly.
 - Per-batch column statistics are recorded in the layout Spark's `SimpleMetricsCachedBatchSerializer`
   expects, so Spark can prune whole cached batches on a predicate before any of them is decoded.
+- The size Spark's planner sees for a cached relation is its decoded Arrow size, not the compressed
+  size it occupies in memory, as with Spark's own cache formats. Compression therefore does not
+  change how queries over a cached relation are planned, such as whether a join broadcasts it.
 
 Relations whose schema Comet's Arrow writer cannot store — interval types, most notably — are
 delegated in full to Spark's default cache format, per relation. Which format a relation uses does
@@ -58,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.
@@ -102,7 +109,7 @@ nowhere to record either that a column is dictionary encoded or the dictionary i
 
 | Config                                                  | Default | Description                                                                                                                                    |
 | ------------------------------------------------------- | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------- |
-| `spark.comet.exec.inMemoryCache.enabled`                | `false` | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup.                                                          |
+| `spark.comet.exec.inMemoryCache.enabled`                | `true`  | Whether to store and scan Spark's in-memory cache in Comet's format. Read at startup.                                                          |
 | `spark.comet.exec.inMemoryCache.compression.codec`      | `zstd`  | Arrow IPC compression codec for cached data: `zstd` or `none`. Affects newly cached data only — a batch records the codec it was written with. |
 | `spark.comet.exec.inMemoryCache.compression.zstd.level` | `1`     | Compression level when the codec is `zstd`. Ignored otherwise.                                                                                 |
 
@@ -190,8 +197,8 @@ format, and the narrower the read, the wider the gap. Measured by the same bench
 | 3 of 6 columns          |                98 ms |               331 ms |     3.4x |
 | 6 of 6 columns          |               410 ms |               623 ms |     1.5x |
 
-This is why the feature is off by default. The cause is not yet established;
-[#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks it.
+This is the main reason the feature is still described as experimental. The cause is not yet
+established; [#5485](https://github.com/apache/datafusion-comet/issues/5485) tracks it.
 
 Comet's serializer exists because Spark's own Arrow cache format
 ([SPARK-57268](https://issues.apache.org/jira/browse/SPARK-57268)) is only available from Spark

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -280,7 +280,7 @@ object CometConf extends ShimCometConf {
         "SparkContext, otherwise caching fails as soon as a block is serialized, including " +
         "the disk half of the default MEMORY_AND_DISK storage level.")
       .booleanConf
-      .createWithDefault(false)
+      .createWithDefault(true)
 
   val COMET_EXEC_IN_MEMORY_CACHE_COMPRESSION_CODEC: ConfigEntry[String] =
     conf("spark.comet.exec.inMemoryCache.compression.codec")

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -386,8 +386,12 @@ 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 only consulted once the native scan is otherwise
+        // possible. See CometInMemoryTableScanExec.recordsObservedMetrics.
+        val nativeScan = nativeCacheEnabled && cometCacheFormat &&
+          !CometInMemoryTableScanExec.recordsObservedMetrics(scan.relation)
 
-        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,14 +402,20 @@ 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}")
             withFallbackReason(
               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) {
+            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,

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)
+  }
+
 }

spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowCachedBatchSerializer.scala

@@ -52,10 +52,14 @@ import org.apache.comet.vector.NativeUtil
  * and length, so `CachedBatchIpc.Projection.load` copies out just the selected columns' byte
  * ranges. The cache manager still owns storage and eviction; this class only changes the cached
  * payload.
+ *
+ * `sizeInBytes` is not the payload's size. It is inherited from `SimpleMetricsCachedBatch`, which
+ * sums the per-column sizes in `stats`, and those are decoded sizes. Spark's planner reads it as
+ * the size of a materialized cached relation, as it does for Spark's own cache formats, which
+ * also report decoded sizes. See `statsRow`.
  */
 private case class CometCachedBatch(
     override val numRows: Int,
-    override val sizeInBytes: Long,
     override val stats: InternalRow,
     bytes: ChunkedByteBuffer)
     extends SimpleMetricsCachedBatch
@@ -350,11 +354,13 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
       values(base + 1) = upper(c)
       values(base + 2) = nulls(c)
       values(base + 3) = numRows
-      // The stored size of the column's own Arrow buffers, taken from the message's buffer
-      // layout, so it is exact rather than an estimate. Cache pruning uses
-      // bounds/null-count/row-count rather than this field, but Spark reserves it and reports it,
-      // so record the real value. The per-batch message framing is not attributed to any column,
-      // so these sum to slightly less than sizeInBytes.
+      // The column's decoded size: the plain length of its own Arrow buffers before compression,
+      // which is what Spark's own Arrow cache format records here. SimpleMetricsCachedBatch sums
+      // these into sizeInBytes, which Spark's planner reads as the size of a materialized cached
+      // relation, for the broadcast threshold and the shuffled hash join build side among others.
+      // The compressed payload can be several times smaller, and a relation reported at that size
+      // would be planned differently from the same relation in Spark's cache format, for example
+      // broadcast where Spark's cache would have it shuffled.
       values(base + 4) = columnSizes(c)
       c += 1
     }
@@ -423,7 +429,6 @@ class ArrowCachedBatchSerializer extends SimpleMetricsCachedBatchSerializer {
 
       CometCachedBatch(
         numRows = numRows,
-        sizeInBytes = bytes.size,
         stats = statsRow(lower, upper, nulls, numRows, columnSizes),
         bytes = bytes)
     }

spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala

@@ -186,9 +186,10 @@ private[comet] object CachedBatchIpc {
    * batch arrives at whatever size the plan above produced. Chunks are appended rather than grown
    * and recopied, so the write also never holds the payload twice.
    *
-   * Returns the message and the on-body compressed size of each top-level column, which the
-   * caller records in the statistics row. The sizes come from the message's own buffer layout, so
-   * they are the real stored sizes rather than an estimate.
+   * Returns the message and the decoded size of each top-level column, which the caller records
+   * in the statistics row. Each size is measured on the batch before compression, from the plain
+   * lengths of the column's own buffers, which is what `getBufferSize` reports for a vector and
+   * what Spark's own Arrow cache format records.
    *
    * Dictionary-encoded columns are decoded to their plain form first. A payload with no Schema
    * message cannot describe a dictionary encoding, and the schema the reader rebuilds from Spark
@@ -215,16 +216,18 @@ private[comet] object CachedBatchIpc {
 
       // Unloaded plain and compressed afterwards rather than by handing the codec to the unloader;
       // see compressed for why.
+      val fields = vectors.map(_.getField)
       val unloader = new VectorUnloader(root, true, NoCompressionCodec.INSTANCE, true)
       val plainBatch = unloader.getRecordBatch
-      val recordBatch =
-        try compressed(plainBatch, codec, allocator)
-        finally plainBatch.close()
+      val (sizes, recordBatch) =
+        try {
+          (columnSizes(fields, plainBatch), compressed(plainBatch, codec, allocator))
+        } finally {
+          plainBatch.close()
+        }
       try {
-        val fields = vectors.map(_.getField)
         // Leaves the batch in the state serializeBatches leaves one. The record batch holds its
-        // own buffers by now, so this does not touch it, and getField still answers afterwards:
-        // clearing releases buffers, not the schema.
+        // own buffers by now, so this does not touch it.
         root.clear()
 
         val out = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate)
@@ -234,7 +237,7 @@ private[comet] object CachedBatchIpc {
         } finally {
           out.close()
         }
-        (out.toChunkedByteBuffer, columnSizes(fields, recordBatch))
+        (out.toChunkedByteBuffer, sizes)
       } finally {
         recordBatch.close()
       }
@@ -608,11 +611,10 @@ private[comet] object CachedBatchIpc {
   }
 
   /**
-   * The on-body compressed size of each top-level column.
+   * The size of each top-level column's buffers in `recordBatch`.
    *
-   * Each column owns the run of buffers its subtree occupies, so its stored size is the sum of
-   * those buffers' recorded lengths. With one payload per batch these are the only per-column
-   * sizes available -- there is no separate stream to measure -- and they are exact.
+   * Each column owns the run of buffers its subtree occupies, so its size is the sum of those
+   * buffers' recorded lengths.
    */
   private def columnSizes(fields: Seq[Field], recordBatch: ArrowRecordBatch): Array[Long] = {
     val buffers = recordBatch.getBuffersLayout