[SPARK-57512][SQL] Materialize surviving RuntimeReplaceable for cached-batch pruning

Address review: a surviving RuntimeReplaceable in a cached scan's pushed-down
predicates was not recognized by CachedBatchSerializer.buildFilter, silently
disabling cached-batch pruning under AQE (the scan is wrapped in a leaf
TableCacheQueryStageExec that the stage-finalization MaterializeRuntimeReplaceable
cannot reach).

Fix it at the pushdown consumer: unfold RuntimeReplaceable in InMemoryTableScanExec
before calling buildFilter, rather than extending the codegen-prep rule into the
non-codegen leaf scan. This covers AQE and non-AQE uniformly and keeps the readable
expression in the plan/EXPLAIN. Document why codegen-materialization lives in
preparations/postStageCreationRules and why the InMemory branch deliberately skips it.

Co-authored-by: Isaac

Author: cloud-fan

sql/core/src/main/scala/org/apache/spark/sql/execution/MaterializeRuntimeReplaceable.scala

@@ -27,10 +27,25 @@ import org.apache.spark.sql.catalyst.trees.TreePattern.RUNTIME_REPLACEABLE
  *
  * A `RuntimeReplaceable` with `eagerReplace = false` is intentionally kept in the plan by the
  * optimizer (see `ReplaceExpressions`) so that a native engine can match the high-level expression
- * directly. This rule then materializes the replacement for the Spark execution path, so Spark
- * codegen/interpreted evaluation behaves exactly as today. It is placed after the columnar/native
- * conversion and before `CollapseCodegenStages`, so a native engine sees the original expression
- * while Spark whole-stage codegen never sees a `RuntimeReplaceable`.
+ * directly. This rule then materializes the replacement for the Spark execution path. It is placed
+ * after the columnar/native conversion and before `CollapseCodegenStages`, so a native engine sees
+ * the original expression while Spark whole-stage codegen never sees a `RuntimeReplaceable`.
+ *
+ * Materializing before codegen is a correctness requirement, not just cleanup. A surviving
+ * `RuntimeReplaceable` evaluates correctly on its own (`eval`/`doGenCode` delegate to
+ * `replacement`), but whole-stage codegen reasons about `references`, input materialization, and
+ * subexpression elimination via the node's `children`, while the emitted code comes from
+ * `replacement`. When `replacement` reads an input differently from `children` (e.g. an
+ * un-simplified branch that reads a column more than once), that mismatch produces invalid
+ * generated code. Eager replacement avoids this because the unfolded form is then simplified by the
+ * optimizer; a survivor's `replacement` is not, so it must be unfolded before codegen.
+ *
+ * This runs wherever a physical plan is finalized into a codegen-bearing form:
+ * `QueryExecution.preparations` (non-AQE) and `AdaptiveSparkPlanExec.postStageCreationRules` (AQE,
+ * applied to every codegen-producing stage). A `RuntimeReplaceable` that only feeds a structural,
+ * non-codegen consumer is materialized at that consumer instead -- see the cached-batch pruning
+ * predicates in `InMemoryTableScanExec`, whose leaf scan never reaches codegen and is unreachable
+ * from AQE stage finalization.
  */
 object MaterializeRuntimeReplaceable extends Rule[SparkPlan] {
   override def apply(plan: SparkPlan): SparkPlan = plan.transformUpWithSubqueries {

sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AdaptiveSparkPlanExec.scala

@@ -713,7 +713,10 @@ case class AdaptiveSparkPlanExec(
       case i: InMemoryTableScanLike =>
         // Apply `queryStageOptimizerRules` so that we can reuse subquery.
         // No need to apply `postStageCreationRules` for `InMemoryTableScanLike`
-        // as it's a leaf node.
+        // as it's a leaf node. In particular, `MaterializeRuntimeReplaceable` is intentionally not
+        // applied here: this scan does not reach whole-stage codegen, and its only expressions that
+        // may hold a surviving `RuntimeReplaceable` are the pushed-down `predicates`, which are
+        // materialized at their consumer in `InMemoryTableScanExec` (see `filteredCachedBatches`).
         val newPlan = optimizeQueryStage(i, isFinalStage = false)
         if (!newPlan.isInstanceOf[InMemoryTableScanLike]) {
           throw SparkException.internalError(

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala

@@ -141,7 +141,18 @@ case class InMemoryTableScanExec(
     val buffers = relation.cacheBuilder.cachedColumnBuffers
 
     if (inMemoryPartitionPruningEnabled) {
-      val filterFunc = relation.cacheBuilder.serializer.buildFilter(predicates, relation.output)
+      // `predicates` may contain a surviving `RuntimeReplaceable` (`eagerReplace = false`), which
+      // is intentionally kept in the plan. `buildFilter` matches on expression shape to build the
+      // cached-batch pruning filter, so it must see the materialized form. We unfold here, at the
+      // consumer, rather than relying on the codegen-prep materialization rule
+      // (`MaterializeRuntimeReplaceable` in `QueryExecution.preparations` /
+      // `AdaptiveSparkPlanExec.postStageCreationRules`): this scan is a leaf that never reaches
+      // whole-stage codegen, and under AQE it is wrapped in a `TableCacheQueryStageExec` that the
+      // stage-finalization rules cannot descend into. Unfolding here covers both the AQE and
+      // non-AQE paths uniformly while keeping the readable expression in the plan/EXPLAIN output.
+      val materializedPredicates = predicates.map(RuntimeReplaceable.unfold)
+      val filterFunc =
+        relation.cacheBuilder.serializer.buildFilter(materializedPredicates, relation.output)
       buffers.mapPartitionsWithIndexInternal(filterFunc)
     } else {
       buffers

sql/core/src/test/scala/org/apache/spark/sql/execution/MaterializeRuntimeReplaceableSuite.scala

@@ -18,11 +18,12 @@
 package org.apache.spark.sql.execution
 
 import org.apache.spark.sql.{QueryTest, Row}
-import org.apache.spark.sql.catalyst.expressions.{Add, Expression, Literal, RuntimeReplaceable}
+import org.apache.spark.sql.catalyst.expressions.{Add, Expression, GreaterThan, Literal, RuntimeReplaceable}
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.BinaryLike
-import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec}
+import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
 
@@ -54,6 +55,34 @@ object WrapAddWithRuntimeReplaceable extends Rule[LogicalPlan] {
   }
 }
 
+/**
+ * A test-only predicate [[RuntimeReplaceable]] (`eagerReplace = false`) whose `replacement` is a
+ * plain [[GreaterThan]] -- a shape that `CachedBatchSerializer.buildFilter` recognizes for
+ * cached-batch pruning. Used to verify that a surviving predicate is materialized at the pruning
+ * consumer (`InMemoryTableScanExec`), so pruning still kicks in.
+ */
+case class TestPredicateRuntimeReplaceable(left: Expression, right: Expression)
+  extends RuntimeReplaceable with BinaryLike[Expression] {
+
+  override lazy val replacement: Expression = GreaterThan(left, right)
+
+  override def eagerReplace: Boolean = false
+
+  override protected def withNewChildrenInternal(
+      newLeft: Expression, newRight: Expression): TestPredicateRuntimeReplaceable =
+    copy(left = newLeft, right = newRight)
+}
+
+/**
+ * Wraps `x > 88` into a surviving [[TestPredicateRuntimeReplaceable]], after `ReplaceExpressions`.
+ */
+object WrapGreaterThanWithRuntimeReplaceable extends Rule[LogicalPlan] {
+  override def apply(plan: LogicalPlan): LogicalPlan = plan.transformAllExpressions {
+    case g: GreaterThan if g.right == Literal(88) =>
+      TestPredicateRuntimeReplaceable(g.left, g.right)
+  }
+}
+
 class MaterializeRuntimeReplaceableSuite extends QueryTest with SharedSparkSession {
 
   private def withExtraOptimization(rule: Rule[LogicalPlan])(f: => Unit): Unit = {
@@ -107,6 +136,63 @@ class MaterializeRuntimeReplaceableSuite extends QueryTest with SharedSparkSessi
     }
   }
 
+  test("SPARK-57512: a surviving RuntimeReplaceable in cached-scan predicates is materialized " +
+      "for partition pruning under AQE") {
+    // Find every InMemoryTableScanExec, descending through AQE query stages (the cached scan is
+    // wrapped in a TableCacheQueryStageExec, which is a leaf to a normal tree walk).
+    def findScans(p: SparkPlan): Seq[InMemoryTableScanExec] = p match {
+      case in: InMemoryTableScanExec => Seq(in)
+      case q: QueryStageExec => findScans(q.plan)
+      case a: AdaptiveSparkPlanExec => findScans(a.executedPlan)
+      case other => other.children.flatMap(findScans)
+    }
+
+    withSQLConf(
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+        // Keep the range partitions intact so pruning is observable.
+        SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false",
+        SQLConf.COLUMN_BATCH_SIZE.key -> "10",
+        SQLConf.IN_MEMORY_PARTITION_PRUNING.key -> "true",
+        SQLConf.IN_MEMORY_TABLE_SCAN_STATISTICS_ENABLED.key -> "true") {
+      withExtraOptimization(WrapGreaterThanWithRuntimeReplaceable) {
+        import testImplicits._
+        // `repartitionByRange` adds a shuffle, so the cached plan is adaptive and its scan is
+        // wrapped in a `TableCacheQueryStageExec` (a leaf the AQE stage-finalization rules cannot
+        // descend into) -- this is the case where the predicate is NOT materialized in the plan.
+        // Range partitioning also keeps values clustered, so batch stats enable pruning.
+        // 100 values, batch size 10 => 10 batches total across 5 partitions.
+        val cached = sparkContext.makeRDD(1 to 100, 5).toDF("key").repartitionByRange(5, $"key")
+        cached.cache()
+        try {
+          // `key > 88` is rewritten into a surviving `TestPredicateRuntimeReplaceable` and pushed
+          // into the cached scan's `predicates`.
+          val df = cached.filter("key > 88")
+          checkAnswer(df, (89 to 100).map(Row(_)))
+
+          val scans = findScans(df.queryExecution.executedPlan)
+          assert(scans.size == 1, s"expected one cached scan, found ${scans.size}")
+          val scan = scans.head
+
+          // The scan is a leaf query stage, so its predicate is not materialized in the plan: the
+          // surviving RuntimeReplaceable is still present. This is exactly why the consumer-side
+          // unfold in `filteredCachedBatches` is needed.
+          assert(
+            scan.predicates.exists(_.exists(_.isInstanceOf[RuntimeReplaceable])),
+            s"Expected a surviving RuntimeReplaceable in the cached scan predicates:\n$scan")
+
+          // Pruning kicked in: fewer than all 10 batches / 5 partitions are read. Without unfolding
+          // the predicate, `buildFilter` would not recognize it and would scan everything.
+          assert(scan.readBatches.value < 10,
+            s"Expected pruning (< 10 batches read), got ${scan.readBatches.value}")
+          assert(scan.readPartitions.value < 5,
+            s"Expected pruning (< 5 partitions read), got ${scan.readPartitions.value}")
+        } finally {
+          cached.unpersist()
+        }
+      }
+    }
+  }
+
   test("a surviving RuntimeReplaceable self-evaluates via its replacement") {
     // `eval` delegates to `replacement` as a backstop for paths that bypass
     // `MaterializeRuntimeReplaceable`.