[SPARK-58208][SQL] Deep-copy stateful expressions before optimization

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

@@ -308,7 +308,7 @@ case class Uniform(
 
   override def withNewChildrenInternal(
       newFirst: Expression, newSecond: Expression, newThird: Expression): Expression =
-    Uniform(newFirst, newSecond, newThird, hideSeed)
+    copy(min = newFirst, max = newSecond, seedExpression = newThird)
 
   override def replacement: Expression = {
     if (Seq(min, max, seedExpression).exists(_.dataType == NullType)) {

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

@@ -2785,9 +2785,11 @@ object ConvertToLocalRelation extends Rule[LogicalPlan] {
     _.containsPattern(LOCAL_RELATION), ruleId) {
     case Project(projectList, LocalRelation(output, data, isStreaming, stream))
         if !projectList.exists(hasUnevaluableExpr) =>
-      val projection = new InterpretedMutableProjection(projectList, output)
+      val freshProjectList = projectList.map(
+        _.freshCopyIfContainsStatefulExpression().asInstanceOf[NamedExpression])
+      val projection = new InterpretedMutableProjection(freshProjectList, output)
       projection.initialize(0)
-      LocalRelation(projectList.map(_.toAttribute), data.map(projection(_).copy()),
+      LocalRelation(freshProjectList.map(_.toAttribute), data.map(projection(_).copy()),
         isStreaming, stream)
 
     case Limit(IntegerLiteral(limit), LocalRelation(output, data, isStreaming, stream)) =>
@@ -2798,7 +2800,8 @@ object ConvertToLocalRelation extends Rule[LogicalPlan] {
 
     case Filter(condition, LocalRelation(output, data, isStreaming, stream))
         if !hasUnevaluableExpr(condition) =>
-      val predicate = Predicate.create(condition, output)
+      val freshCondition = condition.freshCopyIfContainsStatefulExpression()
+      val predicate = Predicate.create(freshCondition, output)
       predicate.initialize(0)
       LocalRelation(output, data.filter(row => predicate.eval(row)), isStreaming, stream)
   }

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

@@ -258,13 +258,28 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]]
    * query operator based on the mapped expressions.
    */
   def mapExpressions(f: Expression => Expression): this.type = {
+    mapExpressions(f, useFastEquals = true)
+  }
+
+  /**
+   * A variant of [[mapExpressions]] that retains structurally equal replacement expressions.
+   */
+  private[sql] def mapExpressionsWithReferenceEquality(
+      f: Expression => Expression): this.type = {
+    mapExpressions(f, useFastEquals = false)
+  }
+
+  private def mapExpressions(
+      f: Expression => Expression,
+      useFastEquals: Boolean): this.type = {
     var changed = false
 
     @inline def transformExpression(e: Expression): Expression = {
       val newE = CurrentOrigin.withOrigin(e.origin) {
         f(e)
       }
-      if (newE.fastEquals(e)) {
+      val unchanged = if (useFastEquals) newE.fastEquals(e) else newE.eq(e)
+      if (unchanged) {
         e
       } else {
         changed = true
@@ -577,6 +592,30 @@ abstract class QueryPlan[PlanType <: QueryPlan[PlanType]]
     transformDownWithSubqueriesAndPruning(AlwaysProcess.fn, UnknownRuleId)(f)
   }
 
+  /**
+   * A variant of [[transformDownWithSubqueries]] that retains structurally equal replacement
+   * plans and expressions.
+   */
+  private[sql] def transformDownWithSubqueriesAndReferenceEquality(
+      f: PartialFunction[PlanType, PlanType]): PlanType = {
+    val g: PartialFunction[PlanType, PlanType] = new PartialFunction[PlanType, PlanType] {
+      override def isDefinedAt(x: PlanType): Boolean = true
+
+      override def apply(plan: PlanType): PlanType = {
+        val transformed = f.applyOrElse[PlanType, PlanType](plan, identity)
+        transformed.mapExpressionsWithReferenceEquality(
+          _.transformDownWithReferenceEquality {
+            case planExpression: PlanExpression[PlanType @unchecked] =>
+              val newPlan = planExpression.plan
+                .transformDownWithSubqueriesAndReferenceEquality(f)
+              planExpression.withNewPlan(newPlan)
+          })
+      }
+    }
+
+    transformDownWithReferenceEquality(g)
+  }
+
   /**
    * Same as `transformUpWithSubqueries` except allows for pruning opportunities.
    */

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala

@@ -471,6 +471,22 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
     transformDownWithPruning(AlwaysProcess.fn, UnknownRuleId)(rule)
   }
 
+  /**
+   * A variant of [[transformDown]] that retains structurally equal replacement nodes.
+   */
+  private[sql] def transformDownWithReferenceEquality(
+      rule: PartialFunction[BaseType, BaseType]): BaseType = {
+    val afterRule = CurrentOrigin.withOrigin(origin) {
+      rule.applyOrElse(this, identity[BaseType])
+    }
+    if (this eq afterRule) {
+      mapChildrenWithReferenceEquality(_.transformDownWithReferenceEquality(rule))
+    } else {
+      afterRule.copyTagsFrom(this)
+      afterRule.mapChildrenWithReferenceEquality(_.transformDownWithReferenceEquality(rule))
+    }
+  }
+
   /**
    * Returns a copy of this node where `rule` has been recursively applied to it and all of its
    * children (pre-order). When `rule` does not apply to a given node it is left unchanged.
@@ -736,6 +752,22 @@ abstract class TreeNode[BaseType <: TreeNode[BaseType]]
     }
   }
 
+  private[sql] final def mapChildrenWithReferenceEquality(
+      f: BaseType => BaseType): BaseType = {
+    val newChildren = children.map(f)
+    if (children.iterator.zip(newChildren.iterator).forall { case (oldChild, newChild) =>
+        oldChild eq newChild
+      }) {
+      this
+    } else {
+      CurrentOrigin.withOrigin(origin) {
+        val res = withNewChildrenInternal(asIndexedSeq(newChildren))
+        res.copyTagsFrom(this)
+        res
+      }
+    }
+  }
+
   /**
    * Args to the constructor that should be copied, but not transformed.
    * These are appended to the transformed args automatically by makeCopy

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala

@@ -64,4 +64,15 @@ class RandomSuite extends SparkFunSuite with ExpressionEvalHelper {
     testUniform(10.0F, 20.0F, 17.604954F)
     testUniform(10L, 20.0F, 17.604954F)
   }
+
+  test("SPARK-58208: Uniform preserves its time zone when copied") {
+    val uniform = Uniform(
+      Literal(10), Literal(20), Literal(0), hideSeed = false, timeZoneId = Some("UTC"))
+    assert(uniform.resolved)
+
+    val copied = uniform.freshCopyIfContainsStatefulExpression().asInstanceOf[Uniform]
+    assert(copied ne uniform)
+    assert(copied.timeZoneId == uniform.timeZoneId)
+    assert(copied.resolved)
+  }
 }

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

@@ -21,12 +21,12 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow, LessThan, Literal, UnaryExpression}
+import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ArrayTransform, Expression, GenericInternalRow, LambdaFunction, LessThan, Literal, NamedLambdaVariable, UnaryExpression}
 import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode}
 import org.apache.spark.sql.catalyst.plans.PlanTest
-import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
+import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan, Project}
 import org.apache.spark.sql.catalyst.rules.RuleExecutor
-import org.apache.spark.sql.types.{DataType, StructType}
+import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, StructType}
 
 
 class ConvertToLocalRelationSuite extends PlanTest {
@@ -87,6 +87,18 @@ class ConvertToLocalRelationSuite extends PlanTest {
 
     comparePlans(optimized, correctAnswer)
   }
+
+  test("SPARK-58208: ConvertToLocalRelation uses fresh stateful project expressions") {
+    val element = NamedLambdaVariable("x", IntegerType, nullable = false)
+    val transform = ArrayTransform(
+      Literal.create(Seq(1, 2), ArrayType(IntegerType, containsNull = false)),
+      LambdaFunction(Add(element, Literal(1)), Seq(element)))
+    val project = Project(Seq(Alias(transform, "v")()), LocalRelation(Nil, Seq(InternalRow.empty)))
+
+    Optimize.execute(project)
+
+    assert(element.value.get() == null)
+  }
 }
 
 

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

@@ -308,15 +308,22 @@ class QueryExecution(
 
   def assertCommandExecuted(): Unit = commandExecuted
 
+  private def cloneWithFreshStatefulExpressions(plan: LogicalPlan): LogicalPlan = {
+    plan.clone().transformDownWithSubqueriesAndReferenceEquality {
+      case node =>
+        node.mapExpressionsWithReferenceEquality(_.freshCopyIfContainsStatefulExpression())
+    }
+  }
+
   private val lazyOptimizedPlan = LazyTry {
     // We need to materialize the commandExecuted here because optimizedPlan is also tracked under
     // the optimizing phase
     assertCommandExecuted()
     executePhase(QueryPlanningTracker.OPTIMIZATION) {
       // clone the plan to avoid sharing the plan instance between different stages like analyzing,
       // optimizing and planning.
-      val plan =
-        sparkSession.sessionState.optimizer.executeAndTrack(withCachedData.clone(), tracker)
+      val plan = sparkSession.sessionState.optimizer.executeAndTrack(
+        cloneWithFreshStatefulExpressions(withCachedData), tracker)
       // We do not want optimized plans to be re-analyzed as literals that have been constant
       // folded and such can cause issues during analysis. While `clone` should maintain the
       // `analyzed` state of the LogicalPlan, we set the plan as analyzed here as well out of

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

@@ -25,7 +25,7 @@ import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListe
 import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, FastOperator, SaveMode}
 import org.apache.spark.sql.catalyst.{QueryPlanningTracker, QueryPlanningTrackerCallback, TableIdentifier}
 import org.apache.spark.sql.catalyst.analysis.{CurrentNamespace, UnresolvedFunction, UnresolvedRelation}
-import org.apache.spark.sql.catalyst.expressions.{Alias, UnsafeRow}
+import org.apache.spark.sql.catalyst.expressions.{Alias, NamedLambdaVariable, RegExpReplace, UnsafeRow}
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.logical.{CommandResult, LogicalPlan, OneRowRelation, Project, ShowTables, SubqueryAlias}
 import org.apache.spark.sql.catalyst.trees.TreeNodeTag
@@ -55,6 +55,22 @@ class QueryExecutionSuite extends SharedSparkSession {
   override protected def sparkConf =
     super.sparkConf.set(SQLConf.ADAPTIVE_MAX_SHUFFLE_HASH_JOIN_LOCAL_MAP_THRESHOLD.key, "0")
 
+  private def collectLambdaVariables(plan: LogicalPlan): Seq[NamedLambdaVariable] = {
+    plan.collect {
+      case node => node.expressions.flatMap(_.collect {
+        case variable: NamedLambdaVariable => variable
+      })
+    }.flatten
+  }
+
+  private def collectRegExpReplaceExpressions(plan: LogicalPlan): Seq[RegExpReplace] = {
+    plan.collect {
+      case node => node.expressions.flatMap(_.collect {
+        case expression: RegExpReplace => expression
+      })
+    }.flatten
+  }
+
   def checkDumpedPlans(path: String, expected: Int): Unit = Utils.tryWithResource(
     Source.fromFile(path)) { source =>
     assert(source.getLines().toList
@@ -105,6 +121,34 @@ class QueryExecutionSuite extends SharedSparkSession {
     }
   }
 
+  test("SPARK-58208: optimizedPlan uses fresh stateful expressions") {
+    val df = spark.range(1).selectExpr("transform(array(id), x -> x + 1) AS v")
+    val queryExecution = df.queryExecution
+
+    val beforeOptimize = collectLambdaVariables(queryExecution.withCachedData)
+    val optimized = collectLambdaVariables(queryExecution.optimizedPlan)
+
+    assert(beforeOptimize.nonEmpty)
+    assert(beforeOptimize.size == optimized.size)
+    beforeOptimize.zip(optimized).foreach { case (before, after) =>
+      assert(before.exprId == after.exprId)
+      assert(before.value ne after.value)
+    }
+  }
+
+  test("SPARK-58208: optimizedPlan keeps structurally equal fresh stateful expressions") {
+    val df = spark.range(1).selectExpr(
+      "regexp_replace(cast(id AS STRING), cast(id AS STRING), 'x') AS v")
+    val queryExecution = df.queryExecution
+
+    val beforeOptimize = collectRegExpReplaceExpressions(queryExecution.withCachedData)
+    val optimized = collectRegExpReplaceExpressions(queryExecution.optimizedPlan)
+
+    assert(beforeOptimize.size == 1)
+    assert(optimized.size == 1)
+    assert(beforeOptimize.head ne optimized.head)
+  }
+
   test("dumping query execution info by invalid path") {
     val path = "1234567890://plans.txt"
     val exception = intercept[IllegalArgumentException] {