[SPARK-57512][SQL] Make RuntimeReplaceable survival opt-in via eagerReplace

Introduce `RuntimeReplaceable.eagerReplace` (default `true`) so existing
expressions keep being rewritten eagerly by `ReplaceExpressions`, while an
expression can opt into surviving into the physical plan by overriding it to
`false`. This enables incremental migration of expressions.

`ReplaceExpressions` rewrites eagerly when `eagerReplace` is true or when the
replacement cannot survive (non-deterministic or unevaluable); otherwise the
node survives and is materialized into its replacement by
`MaterializeRuntimeReplaceable` right before codegen.

Add `MaterializeRuntimeReplaceableSuite` proving an optimizer rule can insert
a surviving `RuntimeReplaceable` that executes correctly.

Co-authored-by: Isaac

Author: cloud-fan

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

@@ -457,8 +457,9 @@ trait RuntimeReplaceable extends Expression {
   // 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 now that a deterministic `RuntimeReplaceable` may
-  // survive into the physical plan: the survival decision relies on an accurate determinism signal.
+  // 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)`,
@@ -471,12 +472,23 @@ trait RuntimeReplaceable extends Expression {
   // are semantically equal.
   override lazy val canonicalized: Expression = replacement.canonicalized
 
+  // 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, 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.
+  // `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 = {

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

@@ -58,17 +58,18 @@ object ReplaceExpressions extends Rule[LogicalPlan] {
 
     case r: RuntimeReplaceable =>
       val replaced = replace(r.replacement)
-      // Only a deterministic, fully-evaluable replacement may survive into the physical plan, to be
-      // matched by a native engine or materialized just before codegen. The other cases must be
-      // rewritten early instead:
+      // 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.
-      // A foldable replacement (e.g. `collation(c1)`) is allowed to survive here; `ConstantFolding`
-      // later materializes it into a literal.
-      if (replaced.deterministic && !replaced.exists(_.isInstanceOf[Unevaluable])) r else replaced
+      val cannotSurvive = !replaced.deterministic || replaced.exists(_.isInstanceOf[Unevaluable])
+      if (r.eagerReplace || cannotSurvive) replaced else r
 
     case _ => e.mapChildren(replace)
   }

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

@@ -21,7 +21,7 @@ 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, Expression, IntegerLiteral, Literal, MapFromEntries, Multiply, NamedExpression, NullIf, Remainder, RuntimeReplaceable}
+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.aggregate.Sum
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LocalRelation, LogicalPlan, OneRowRelation, Project}
@@ -347,20 +347,11 @@ class OptimizerSuite extends PlanTest {
       val plan = Project(Alias(nullIf, "out")() :: Nil, OneRowRelation()).analyze
       val optimized = optimizer.execute(plan)
 
-      // NullIf is now deterministic and evaluable, so it survives ReplaceExpressions and is only
-      // unfolded at the physical layer (MaterializeRuntimeReplaceable). Materialize it here to
-      // inspect the form that reaches execution.
-      def materialize(e: Expression): Expression = e match {
-        case r: RuntimeReplaceable => materialize(r.replacement)
-        case _ => e.mapChildren(materialize)
-      }
-      val materialized = optimized.expressions.map(materialize)
-
-      assert(materialized.exists(_.exists {
+      assert(optimized.expressions.exists(_.exists {
         case Literal(null, BooleanType) => true
         case _ => false
       }))
-      assert(materialized.forall(!_.exists(_.isInstanceOf[RuntimeReplaceable])))
+      assert(optimized.expressions.forall(!_.exists(_.isInstanceOf[RuntimeReplaceable])))
     }
   }
 }

sql/connect/common/src/test/resources/query-tests/explain-results/function_collation.explain

@@ -1,2 +1,2 @@
-Project [collation(g#0) AS collation(g)#0]
+Project [SYSTEM.BUILTIN.UTF8_BINARY AS collation(g)#0]
 +- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]

sql/connect/common/src/test/resources/query-tests/explain-results/function_schema_of_csv.explain

@@ -1,2 +1,2 @@
-Project [schema_of_csv(1|abc, (sep,|)) AS schema_of_csv(1|abc)#0]
+Project [invoke(SchemaOfCsvEvaluator(Map(sep -> |)).evaluate(1|abc)) AS schema_of_csv(1|abc)#0]
 +- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]

sql/connect/common/src/test/resources/query-tests/explain-results/function_schema_of_json.explain

@@ -1,2 +1,2 @@
-Project [schema_of_json([{"col":01}]) AS schema_of_json([{"col":01}])#0]
+Project [invoke(SchemaOfJsonEvaluator(Map()).evaluate([{"col":01}])) AS schema_of_json([{"col":01}])#0]
 +- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]

sql/connect/common/src/test/resources/query-tests/explain-results/function_schema_of_json_with_options.explain

@@ -1,2 +1,2 @@
-Project [schema_of_json([{"col":01}], (allowNumericLeadingZeros,true)) AS schema_of_json([{"col":01}])#0]
+Project [invoke(SchemaOfJsonEvaluator(Map(allowNumericLeadingZeros -> true)).evaluate([{"col":01}])) AS schema_of_json([{"col":01}])#0]
 +- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]

sql/connect/common/src/test/resources/query-tests/explain-results/function_try_to_binary.explain

@@ -1,2 +1,2 @@
-Project [tryeval(to_binary(g#0, Some(format), true)) AS try_to_binary(g, format)#0]
+Project [tryeval(null) AS try_to_binary(g, format)#0]
 +- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]

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

@@ -22,15 +22,15 @@ import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreePattern.RUNTIME_REPLACEABLE
 
 /**
- * Prototype: materializes any [[RuntimeReplaceable]] that survived the logical optimizer into its
- * `replacement`, on the physical plan, right before whole-stage codegen.
+ * Materializes any [[RuntimeReplaceable]] that survived the logical optimizer into its
+ * `replacement`, on the physical plan.
  *
- * In this prototype `RuntimeReplaceable` is intentionally left in the plan by the optimizer (scalar
- * `ReplaceExpressions` is disabled) so that a native engine (Photon) could match the semantic
- * expression directly. This rule then materializes the replacement for the Spark execution path, so
- * Spark codegen/interpreted evaluation behaves exactly as today. Placed after the columnar/native
- * conversion and before `CollapseCodegenStages`, so a native engine sees the origin while Spark
- * sees the replacement.
+ * 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`.
  */
 object MaterializeRuntimeReplaceable extends Rule[SparkPlan] {
   override def apply(plan: SparkPlan): SparkPlan = plan.transformUpWithSubqueries {

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

@@ -775,9 +775,9 @@ object QueryExecution {
       DisableUnnecessaryBucketedScan,
       ApplyColumnarRulesAndInsertTransitions(
         sparkSession.sessionState.columnarRules, outputsColumnar = false),
-      // Prototype: materialize any RuntimeReplaceable that survived the optimizer into its
-      // replacement for the Spark execution path. After columnar/native conversion (so a native
-      // engine sees the origin), before codegen (so Spark codegen never sees a RuntimeReplaceable).
+      // Materialize any RuntimeReplaceable that survived the optimizer into its replacement for the
+      // Spark execution path. After columnar/native conversion (so a native engine sees the original
+      // expression), before codegen (so Spark codegen never sees a RuntimeReplaceable).
       MaterializeRuntimeReplaceable,
       CollapseCodegenStages()) ++
       (if (subquery) {