[SPARK-57512][SQL] Allow RuntimeReplaceable to opt out of eager replacement and survive into the physical plan

Add `RuntimeReplaceable.eagerReplace` (default `true`, preserving current behavior). An
expression can override it to `false` to survive the logical optimizer into the physical
plan, so a native engine can match the high-level expression directly and the optimized
plan stays readable. `ReplaceExpressions` still rewrites eagerly when the replacement
cannot survive (non-deterministic, or contains an `Unevaluable` other than
`AttributeReference`); `RuntimeReplaceableAggregate` is always rewritten.

Supporting changes:
- `eval`/`doGenCode`/`deterministic`/`foldable` delegate to `replacement`; a new
  physical-preparation rule `MaterializeRuntimeReplaceable` unfolds survivors before
  `CollapseCodegenStages` (in both `QueryExecution.preparations` and AQE
  `postStageCreationRules`). This is required for correctness: whole-stage codegen reasons
  about `references`/input materialization/CSE via `children`, while the emitted code comes
  from `replacement`, so an un-materialized survivor can produce invalid generated code.
- `FoldablePropagation` only propagates literals; `NormalizePlan` fully unfolds.
- Structure-interpreting consumers that receive a predicate leaving Spark's evaluation
  engine unfold survivors at their boundary: data source filter pushdown
  (`DataSourceStrategy.translateLeafNodeFilter` for V1, `V2ExpressionBuilder` for V2 which
  also covers aggregate/group-by) and cached-batch pruning
  (`InMemoryTableScanExec.buildFilter`).
- First adopter: `MultiGetJsonObject` (inserted by `OptimizeCsvJsonExprs` after
  `ReplaceExpressions`) becomes a surviving `RuntimeReplaceable` with an `Invoke`
  replacement.

Co-authored-by: Isaac

Author: cloud-fan

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/Expression.scala

@@ -454,19 +454,59 @@ trait RuntimeReplaceable extends Expression {
   override val nodePatterns: Seq[TreePattern] = Seq(RUNTIME_REPLACEABLE)
   override def nullable: Boolean = replacement.nullable
   override def dataType: DataType = replacement.dataType
+  // The actual evaluation is delegated to `replacement`, so determinism must reflect `replacement`,
+  // not this expression's `children` (which are the original arguments). For example, the children
+  // of `Uniform` are literal bounds and a seed (all deterministic), while its `replacement` is a
+  // non-deterministic `Rand`. This matters once a `RuntimeReplaceable` may survive into the
+  // physical plan (see `eagerReplace`): the survival decision relies on an accurate determinism
+  // signal.
+  override lazy val deterministic: Boolean = replacement.deterministic
+  // Foldability is also derived from `replacement` rather than this expression's `children`. Note
+  // that this can yield a foldable expression that still has references (e.g. `collation(c1)`,
+  // whose value depends only on the child's type): such an expression is materialized into a
+  // literal by `ConstantFolding`, and `FoldablePropagation` only propagates literals, never bare
+  // foldables.
+  override def foldable: Boolean = replacement.foldable
   // As this expression gets replaced at optimization with its `child" expression,
   // two `RuntimeReplaceable` are considered to be semantically equal if their "child" expressions
   // are semantically equal.
   override lazy val canonicalized: Expression = replacement.canonicalized
 
-  final override def eval(input: InternalRow = null): Any = {
-    // For convenience, we allow to evaluate `RuntimeReplaceable` expressions, in case we need to
-    // get a constant from foldable expression before the query execution starts.
-    assert(input == null)
-    replacement.eval()
+  // Whether `ReplaceExpressions` should rewrite this expression into its `replacement` eagerly, in
+  // the logical optimizer. This is `true` by default, which preserves the historical behavior where
+  // a `RuntimeReplaceable` never reaches the physical plan. An expression can override this to
+  // `false` to survive into the physical plan (e.g. so a native engine can match the high-level
+  // expression directly); such a survivor is materialized into its `replacement` right before
+  // codegen by `MaterializeRuntimeReplaceable`. Note that an expression that opts out can still be
+  // rewritten eagerly if its `replacement` cannot survive (non-deterministic or unevaluable); see
+  // `ReplaceExpressions`.
+  def eagerReplace: Boolean = true
+
+  // `RuntimeReplaceable` expressions are normally rewritten into their `replacement` by the
+  // `ReplaceExpressions` rule before execution. However, an expression with `eagerReplace = false`
+  // survives into the physical plan, and a `RuntimeReplaceable` may also be produced *after*
+  // `ReplaceExpressions` has run (e.g. by an optimizer rule). To keep such an expression evaluable
+  // without depending on the rewrite, both `eval` and `doGenCode` delegate to `replacement`. As
+  // `replacement` is derived from this expression's children, it is bound and code-generated
+  // together with them, so the delegation observes the same input row.
+  final override def eval(input: InternalRow = null): Any = replacement.eval(input)
+
+  final override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    val childGen = replacement.genCode(ctx)
+    ev.copy(code = childGen.code, isNull = childGen.isNull, value = childGen.value)
+  }
+}
+
+object RuntimeReplaceable {
+  /**
+   * Fully unfolds every [[RuntimeReplaceable]] in `e` into its `replacement`, recursively. The
+   * result contains no [[RuntimeReplaceable]] node. This is the canonical "materialize" transform
+   * shared by `MaterializeRuntimeReplaceable` (physical plan) and `NormalizePlan` (comparison).
+   */
+  def unfold(e: Expression): Expression = e match {
+    case r: RuntimeReplaceable => unfold(r.replacement)
+    case _ => e.mapChildren(unfold)
   }
-  final override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode =
-    throw QueryExecutionErrors.cannotGenerateCodeForExpressionError(this)
 }
 
 /**

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/interfaces.scala

@@ -206,8 +206,11 @@ case class AggregateExpression(
  */
 abstract class AggregateFunction extends Expression {
 
-  /** An aggregate function is not foldable. */
-  final override def foldable: Boolean = false
+  // An aggregate function is not foldable. This is not `final` so that
+  // `RuntimeReplaceableAggregate` can inherit `RuntimeReplaceable.foldable` (which delegates to
+  // `replacement`); since such a replacement is itself an aggregate, the effective foldability
+  // stays `false`.
+  override def foldable: Boolean = false
 
   /** The schema of the aggregation buffer. */
   def aggBufferSchema: StructType

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/jsonExpressions.scala

@@ -164,6 +164,7 @@ case class MultiGetJsonObject(
     fieldNames: Seq[String],
     fallbackPaths: Seq[String])
   extends UnaryExpression
+  with RuntimeReplaceable
   with ExpectsInputTypes {
 
   require(
@@ -182,36 +183,29 @@ case class MultiGetJsonObject(
 
   override def nullable: Boolean = true
 
-  // This internal unary expression always returns null when its JSON child is null.
-  override def nullIntolerant: Boolean = true
-
   override def prettyName: String = "multi_get_json_object"
 
-  final override val nodePatterns: Seq[TreePattern] = Seq(GET_JSON_OBJECT)
+  final override val nodePatterns: Seq[TreePattern] = Seq(GET_JSON_OBJECT, RUNTIME_REPLACEABLE)
+
+  // This expression is produced by `OptimizeCsvJsonExprs`, which runs after `ReplaceExpressions`,
+  // so it must opt out of eager replacement to survive into the physical plan. It keeps the
+  // readable `multi_get_json_object` node in the optimized plan and is materialized into its
+  // `replacement` just before codegen by `MaterializeRuntimeReplaceable`.
+  override def eagerReplace: Boolean = false
 
   @transient
   private lazy val evaluator = MultiGetJsonObjectEvaluator(
     fieldNames,
     fallbackPaths.map(UTF8String.fromString))
 
-  override def eval(input: InternalRow): Any = {
-    evaluator.evaluate(json.eval(input).asInstanceOf[UTF8String])
-  }
-
-  override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    val refEvaluator = ctx.addReferenceObj("evaluator", evaluator)
-    val jsonEval = json.genCode(ctx)
-    val resultType = CodeGenerator.javaType(dataType)
-    ev.copy(code = code"""
-       |${jsonEval.code}
-       |boolean ${ev.isNull} = ${jsonEval.isNull};
-       |$resultType ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
-       |if (!${ev.isNull}) {
-       |  ${ev.value} = ($resultType) $refEvaluator.evaluate(${jsonEval.value});
-       |  ${ev.isNull} = ${ev.value} == null;
-       |}
-       |""".stripMargin)
-  }
+  // Delegates evaluation to `MultiGetJsonObjectEvaluator`. `Invoke`'s default `propagateNull`
+  // returns null when the JSON child is null, matching the original null-intolerant behavior.
+  override def replacement: Expression = Invoke(
+    Literal.create(evaluator, ObjectType(classOf[MultiGetJsonObjectEvaluator])),
+    "evaluate",
+    dataType,
+    Seq(json),
+    Seq(json.dataType))
 
   override protected def withNewChildInternal(newChild: Expression): MultiGetJsonObject =
     copy(json = newChild)

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala

@@ -228,7 +228,8 @@ case class Uniform(
   def this(min: Expression, max: Expression, seedExpression: Expression) =
     this(min, max, seedExpression, hideSeed = false)
 
-  final override lazy val deterministic: Boolean = false
+  // `deterministic` is inherited from `RuntimeReplaceable`, which delegates to `replacement` (a
+  // non-deterministic `Rand`-based expression unless an argument is null).
   override def nodePatternsInternal(): Seq[TreePattern] =
     Seq(RUNTIME_REPLACEABLE, EXPRESSION_WITH_RANDOM_SEED)
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/expressions.scala

@@ -1118,7 +1118,13 @@ object FoldablePropagation extends Rule[LogicalPlan] {
 
   private def collectFoldables(expressions: Seq[NamedExpression]) = {
     AttributeMap(expressions.collect {
-      case a: Alias if a.child.foldable => (a.toAttribute, a)
+      // Only propagate literals. A foldable expression is not necessarily a self-contained
+      // constant: a `RuntimeReplaceable` such as `collation(c1)` is foldable (its value depends
+      // only on the child's type) yet still references its children, so substituting it for its
+      // alias elsewhere would leave those references dangling. Such expressions are turned into
+      // literals by `ConstantFolding` (in the same fixed-point batch), after which they propagate
+      // safely.
+      case a: Alias if a.child.isInstanceOf[Literal] => (a.toAttribute, a)
     })
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/finishAnalysis.scala

@@ -53,7 +53,28 @@ object ReplaceExpressions extends Rule[LogicalPlan] {
   }
 
   private def replace(e: Expression): Expression = e match {
-    case r: RuntimeReplaceable => replace(r.replacement)
+    // Aggregates can never self-evaluate (no real aggregation buffer), so always rewrite early.
+    case r: RuntimeReplaceableAggregate => replace(r.replacement)
+
+    case r: RuntimeReplaceable =>
+      val replaced = replace(r.replacement)
+      // By default (`eagerReplace = true`) a `RuntimeReplaceable` is rewritten here, so it never
+      // reaches the physical plan. An expression with `eagerReplace = false` is instead allowed to
+      // survive into the physical plan (to be matched by a native engine or materialized just
+      // before codegen). Even then, a survivor must be rewritten early if its replacement cannot
+      // survive:
+      //  - A non-deterministic replacement (e.g. the `Rand` inside `uniform`) carries mutable
+      //    per-partition state that must be initialized before eval. That state cannot be reliably
+      //    initialized through the `lazy val replacement`, which tree transforms may re-create.
+      //  - A replacement that contains an `Unevaluable` expression (e.g. `With`) depends on a later
+      //    logical rule (such as `RewriteWithExpression`) that can only run in the logical phase.
+      //    `AttributeReference` is `Unevaluable` too but is bound at execution like any input
+      //    column, so it does not block survival -- we reuse `ConvertToLocalRelation`'s check,
+      //    which already excludes it.
+      val cannotSurvive =
+        !replaced.deterministic || ConvertToLocalRelation.hasUnevaluableExpr(replaced)
+      if (r.eagerReplace || cannotSurvive) replaced else r
+
     case _ => e.mapChildren(replace)
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/NormalizePlan.scala

@@ -23,9 +23,8 @@ import org.apache.spark.sql.catalyst.analysis.NormalizeableRelation
 import org.apache.spark.sql.catalyst.analysis.resolver.ResolverTag
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
-import org.apache.spark.sql.catalyst.optimizer.ReplaceExpressions
 import org.apache.spark.sql.catalyst.plans.logical._
-import org.apache.spark.sql.catalyst.trees.TreePattern.PLAN_EXPRESSION
+import org.apache.spark.sql.catalyst.trees.TreePattern.{PLAN_EXPRESSION, RUNTIME_REPLACEABLE}
 
 /**
  * Object that handles normalization of operators and expressions. Used when comparing plans.
@@ -78,8 +77,18 @@ object NormalizePlan extends PredicateHelper {
    * [[InheritAnalysisRules]] is the replacement expression, the original expression will be lost
    * and timezone will never be applied. This causes inconsistencies, because fixed-point semantic
    * is to ALWAYS apply timezone, regardless of whether the Cast actually needs it.
+   *
+   * Note: this unconditionally unfolds every [[RuntimeReplaceable]] into its `replacement`. It
+   * intentionally does NOT reuse the `ReplaceExpressions` optimizer rule, which now keeps
+   * deterministic, evaluable [[RuntimeReplaceable]] nodes in the plan (they are materialized later,
+   * before codegen). Normalization still needs the fully-unfolded form so that the non-child
+   * `parameters` of [[InheritAnalysisRules]] (e.g. the original, un-timezoned cast) are dropped.
    */
-  def normalizeRuntimeReplaceable(plan: LogicalPlan): LogicalPlan = ReplaceExpressions(plan)
+  def normalizeRuntimeReplaceable(plan: LogicalPlan): LogicalPlan = {
+    plan.transformWithPruning(_.containsAnyPattern(RUNTIME_REPLACEABLE)) {
+      case p => p.mapExpressions(RuntimeReplaceable.unfold)
+    }
+  }
 
   /**
    * Since attribute references are given globally unique ids during analysis,

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala

@@ -376,6 +376,11 @@ class V2ExpressionBuilder(e: Expression, isPredicate: Boolean = false) extends L
       } else {
         None
       }
+    // A surviving `RuntimeReplaceable` (`eagerReplace = false`) that no explicit case above pushes
+    // natively is a Spark-internal optimizer node the connector cannot understand. Fall back to its
+    // concrete `replacement` so the lowered form can still be pushed -- same boundary rationale as
+    // `DataSourceStrategy.translateLeafNodeFilter` (V1) and `CachedBatchSerializer.buildFilter`.
+    case r: RuntimeReplaceable => generateExpression(r.replacement, isPredicate)
     case _ => None
   }
 

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizerSuite.scala

@@ -21,11 +21,12 @@ import org.apache.spark.SparkException
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ArrayCompact, AttributeReference, CreateArray, CreateStruct, IntegerLiteral, Literal, MapFromEntries, Multiply, NamedExpression, NullIf, Remainder, RuntimeReplaceable}
+import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ArrayCompact, AttributeReference, CreateArray, CreateStruct, Expression, IntegerLiteral, Literal, MapFromEntries, Multiply, NamedExpression, NullIf, Remainder, RuntimeReplaceable}
 import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LocalRelation, LogicalPlan, OneRowRelation, Project}
 import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.UnaryLike
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{ArrayType, BooleanType, IntegerType, MapType, StructField, StructType}
 
@@ -38,6 +39,23 @@ object DecrementLiterals extends Rule[LogicalPlan] {
   }
 }
 
+/**
+ * A test-only [[RuntimeReplaceable]] that opts out of eager replacement and whose `replacement`
+ * references an input column (via an [[AttributeReference]], which is `Unevaluable`). Used to check
+ * that `ReplaceExpressions` lets such a column-based survivor through, rather than treating the
+ * `AttributeReference` as a reason to rewrite eagerly.
+ */
+case class ColumnBasedRuntimeReplaceable(child: Expression)
+  extends RuntimeReplaceable with UnaryLike[Expression] {
+
+  override lazy val replacement: Expression = Add(child, Literal(1))
+
+  override def eagerReplace: Boolean = false
+
+  override protected def withNewChildInternal(
+      newChild: Expression): ColumnBasedRuntimeReplaceable = copy(child = newChild)
+}
+
 class OptimizerSuite extends PlanTest {
   test("Optimizer exceeds max iterations") {
     val iterations = 5
@@ -354,4 +372,24 @@ class OptimizerSuite extends PlanTest {
       assert(optimized.expressions.forall(!_.exists(_.isInstanceOf[RuntimeReplaceable])))
     }
   }
+
+  test("SPARK-57512: a RuntimeReplaceable with a column-based replacement survives " +
+      "ReplaceExpressions") {
+    val optimizer = new SimpleTestOptimizer() {
+      override def defaultBatches: Seq[Batch] =
+        Batch("test", fixedPoint,
+          ReplaceExpressions) :: Nil
+    }
+
+    val relation = LocalRelation($"a".int)
+    val wrapper = ColumnBasedRuntimeReplaceable(relation.output.head)
+    val plan = Project(Alias(wrapper, "out")() :: Nil, relation).analyze
+    val optimized = optimizer.execute(plan)
+
+    // The replacement references an input column (an `AttributeReference`, which is `Unevaluable`),
+    // but the survivor opts out of eager replacement, so `ReplaceExpressions` keeps it in the plan
+    // instead of rewriting it into the bare `Add`.
+    assert(optimized.expressions.exists(_.exists(_.isInstanceOf[ColumnBasedRuntimeReplaceable])),
+      s"Expected the column-based RuntimeReplaceable to survive:\n$optimized")
+  }
 }

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

@@ -0,0 +1,55 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.execution
+
+import org.apache.spark.sql.catalyst.expressions.RuntimeReplaceable
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.trees.TreePattern.RUNTIME_REPLACEABLE
+
+/**
+ * Materializes any [[RuntimeReplaceable]] that survived the logical optimizer into its
+ * `replacement`, on the physical plan.
+ *
+ * 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. 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 {
+    case p if p.expressions.exists(_.containsPattern(RUNTIME_REPLACEABLE)) =>
+      p.mapExpressions(RuntimeReplaceable.unfold)
+  }
+}