[SPARK-41049][SQL] Revisit stateful expression handling

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

@@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.DataTypeMismatch
 import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateFunction
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.expressions.codegen.Block._
-import org.apache.spark.sql.catalyst.trees.{BinaryLike, LeafLike, QuaternaryLike, SQLQueryContext, TernaryLike, TreeNode, UnaryLike}
+import org.apache.spark.sql.catalyst.trees.{BinaryLike, CurrentOrigin, LeafLike, QuaternaryLike, SQLQueryContext, TernaryLike, TreeNode, UnaryLike}
 import org.apache.spark.sql.catalyst.trees.TreePattern.{RUNTIME_REPLACEABLE, TreePattern}
 import org.apache.spark.sql.catalyst.util.truncatedString
 import org.apache.spark.sql.errors.{QueryErrorsBase, QueryExecutionErrors}
@@ -47,8 +47,6 @@ import org.apache.spark.sql.types._
  * There are a few important traits or abstract classes:
  *
  * - [[Nondeterministic]]: an expression that is not deterministic.
- * - [[Stateful]]: an expression that contains mutable state. For example, MonotonicallyIncreasingID
- *                 and Rand. A stateful expression is always non-deterministic.
  * - [[Unevaluable]]: an expression that is not supposed to be evaluated.
  * - [[CodegenFallback]]: an expression that does not have code gen implemented and falls back to
  *                        interpreted mode.
@@ -127,6 +125,54 @@ abstract class Expression extends TreeNode[Expression] {
 
   def references: AttributeSet = _references
 
+  /**
+   * Returns true if the expression contains mutable state.
+   *
+   * A stateful expression should never be evaluated multiple times for a single row. This should
+   * only be a problem for interpreted execution. This can be prevented by creating fresh copies
+   * of the stateful expression before execution. A common example to trigger this issue:
+   * {{{
+   *   val rand = functions.rand()
+   *   df.select(rand, rand) // These 2 rand should not share a state.
+   * }}}
+   */
+  def stateful: Boolean = false
+
+  /**
+   * Returns a copy of this expression where all stateful expressions are replaced with fresh
+   * uninitialized copies. If the expression contains no stateful expressions then the original
+   * expression is returned.
+   */
+  def freshCopyIfContainsStatefulExpression(): Expression = {
+    val childrenIndexedSeq: IndexedSeq[Expression] = children match {
+      case types: IndexedSeq[Expression] => types
+      case other => other.toIndexedSeq
+    }
+    val newChildren = childrenIndexedSeq.map(_.freshCopyIfContainsStatefulExpression())
+    // A more efficient version of `children.zip(newChildren).exists(_ ne _)`
+    val anyChildChanged = {
+      val size = newChildren.length
+      var i = 0
+      var res: Boolean = false
+      while (!res && i < size) {
+        res |= (childrenIndexedSeq(i) ne newChildren(i))
+        i += 1
+      }
+      res
+    }
+    // If the children contain stateful expressions and get copied, or this expression is stateful,
+    // copy this expression with the new children.
+    if (anyChildChanged || stateful) {
+      CurrentOrigin.withOrigin(origin) {
+        val res = withNewChildrenInternal(newChildren)
+        res.copyTagsFrom(this)
+        res
+      }
+    } else {
+      this
+    }
+  }
+
   /** Returns the result of evaluating this expression on a given input Row */
   def eval(input: InternalRow = null): Any
 
@@ -472,33 +518,6 @@ trait ConditionalExpression extends Expression {
   def branchGroups: Seq[Seq[Expression]]
 }
 
-/**
- * An expression that contains mutable state. A stateful expression is always non-deterministic
- * because the results it produces during evaluation are not only dependent on the given input
- * but also on its internal state.
- *
- * The state of the expressions is generally not exposed in the parameter list and this makes
- * comparing stateful expressions problematic because similar stateful expressions (with the same
- * parameter list) but with different internal state will be considered equal. This is especially
- * problematic during tree transformations. In order to counter this the `fastEquals` method for
- * stateful expressions only returns `true` for the same reference.
- *
- * A stateful expression should never be evaluated multiple times for a single row. This should
- * only be a problem for interpreted execution. This can be prevented by creating fresh copies
- * of the stateful expression before execution, these can be made using the `freshCopy` function.
- */
-trait Stateful extends Nondeterministic {
-  /**
-   * Return a fresh uninitialized copy of the stateful expression.
-   */
-  def freshCopy(): Stateful
-
-  /**
-   * Only the same reference is considered equal.
-   */
-  override def fastEquals(other: TreeNode[_]): Boolean = this eq other
-}
-
 /**
  * A leaf expression, i.e. one without any child expressions.
  */

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

@@ -0,0 +1,45 @@
+/*
+ * 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.catalyst.expressions
+
+import org.apache.spark.sql.internal.SQLConf
+
+// A helper class to evaluate expressions.
+trait ExpressionsEvaluator {
+  protected lazy val runtime =
+    new SubExprEvaluationRuntime(SQLConf.get.subexpressionEliminationCacheMaxEntries)
+
+  protected def prepareExpressions(
+      exprs: Seq[Expression],
+      subExprEliminationEnabled: Boolean): Seq[Expression] = {
+    // We need to make sure that we do not reuse stateful expressions.
+    val cleanedExpressions = exprs.map(_.freshCopyIfContainsStatefulExpression())
+    if (subExprEliminationEnabled) {
+      runtime.proxyExpressions(cleanedExpressions)
+    } else {
+      cleanedExpressions
+    }
+  }
+
+  /**
+   * Initializes internal states given the current partition index.
+   * This is used by nondeterministic expressions to set initial states.
+   * The default implementation does nothing.
+   */
+  def initialize(partitionIndex: Int): Unit = {}
+}

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

@@ -36,18 +36,12 @@ class InterpretedMutableProjection(expressions: Seq[Expression]) extends Mutable
     this(bindReferences(expressions, inputSchema))
 
   private[this] val subExprEliminationEnabled = SQLConf.get.subexpressionEliminationEnabled
-  private[this] lazy val runtime =
-    new SubExprEvaluationRuntime(SQLConf.get.subexpressionEliminationCacheMaxEntries)
-  private[this] val exprs = if (subExprEliminationEnabled) {
-    runtime.proxyExpressions(expressions)
-  } else {
-    expressions
-  }
+  private[this] val exprs = prepareExpressions(expressions, subExprEliminationEnabled)
 
   private[this] val buffer = new Array[Any](expressions.size)
 
   override def initialize(partitionIndex: Int): Unit = {
-    expressions.foreach(_.foreach {
+    exprs.foreach(_.foreach {
       case n: Nondeterministic => n.initialize(partitionIndex)
       case _ =>
     })
@@ -117,10 +111,6 @@ object InterpretedMutableProjection {
    * Returns a [[MutableProjection]] for given sequence of bound Expressions.
    */
   def createProjection(exprs: Seq[Expression]): MutableProjection = {
-    // We need to make sure that we do not reuse stateful expressions.
-    val cleanedExpressions = exprs.map(_.transform {
-      case s: Stateful => s.freshCopy()
-    })
-    new InterpretedMutableProjection(cleanedExpressions)
+    new InterpretedMutableProjection(exprs)
   }
 }

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

@@ -32,13 +32,7 @@ import org.apache.spark.sql.types._
 class InterpretedSafeProjection(expressions: Seq[Expression]) extends Projection {
 
   private[this] val subExprEliminationEnabled = SQLConf.get.subexpressionEliminationEnabled
-  private[this] lazy val runtime =
-    new SubExprEvaluationRuntime(SQLConf.get.subexpressionEliminationCacheMaxEntries)
-  private[this] val exprs = if (subExprEliminationEnabled) {
-    runtime.proxyExpressions(expressions)
-  } else {
-    expressions
-  }
+  private[this] val exprs = prepareExpressions(expressions, subExprEliminationEnabled)
 
   private[this] val mutableRow = new SpecificInternalRow(expressions.map(_.dataType))
 
@@ -106,6 +100,13 @@ class InterpretedSafeProjection(expressions: Seq[Expression]) extends Projection
     case _ => identity
   }
 
+  override def initialize(partitionIndex: Int): Unit = {
+    expressions.foreach(_.foreach {
+      case n: Nondeterministic => n.initialize(partitionIndex)
+      case _ =>
+    })
+  }
+
   override def apply(row: InternalRow): InternalRow = {
     if (subExprEliminationEnabled) {
       runtime.setInput(row)
@@ -130,10 +131,6 @@ object InterpretedSafeProjection {
    * Returns an [[SafeProjection]] for given sequence of bound Expressions.
    */
   def createProjection(exprs: Seq[Expression]): Projection = {
-    // We need to make sure that we do not reuse stateful expressions.
-    val cleanedExpressions = exprs.map(_.transform {
-      case s: Stateful => s.freshCopy()
-    })
-    new InterpretedSafeProjection(cleanedExpressions)
+    new InterpretedSafeProjection(exprs)
   }
 }

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

@@ -35,13 +35,7 @@ class InterpretedUnsafeProjection(expressions: Array[Expression]) extends Unsafe
   import InterpretedUnsafeProjection._
 
   private[this] val subExprEliminationEnabled = SQLConf.get.subexpressionEliminationEnabled
-  private[this] lazy val runtime =
-    new SubExprEvaluationRuntime(SQLConf.get.subexpressionEliminationCacheMaxEntries)
-  private[this] val exprs = if (subExprEliminationEnabled) {
-    runtime.proxyExpressions(expressions)
-  } else {
-    expressions.toSeq
-  }
+  private[this] val exprs = prepareExpressions(expressions, subExprEliminationEnabled)
 
   /** Number of (top level) fields in the resulting row. */
   private[this] val numFields = expressions.length
@@ -106,11 +100,7 @@ object InterpretedUnsafeProjection {
    * Returns an [[UnsafeProjection]] for given sequence of bound Expressions.
    */
   def createProjection(exprs: Seq[Expression]): UnsafeProjection = {
-    // We need to make sure that we do not reuse stateful expressions.
-    val cleanedExpressions = exprs.map(_.transform {
-      case s: Stateful => s.freshCopy()
-    })
-    new InterpretedUnsafeProjection(cleanedExpressions.toArray)
+    new InterpretedUnsafeProjection(exprs.toArray)
   }
 
   /**

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

@@ -48,7 +48,7 @@ import org.apache.spark.sql.types.{DataType, LongType}
   """,
   since = "1.4.0",
   group = "misc_funcs")
-case class MonotonicallyIncreasingID() extends LeafExpression with Stateful {
+case class MonotonicallyIncreasingID() extends LeafExpression with Nondeterministic {
 
   /**
    * Record ID within each partition. By being transient, count's value is reset to 0 every time
@@ -58,11 +58,17 @@ case class MonotonicallyIncreasingID() extends LeafExpression with Stateful {
 
   @transient private[this] var partitionMask: Long = _
 
+  override def stateful: Boolean = true
+
   override protected def initializeInternal(partitionIndex: Int): Unit = {
     count = 0L
     partitionMask = partitionIndex.toLong << 33
   }
 
+  override def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = {
+    MonotonicallyIncreasingID()
+  }
+
   override def nullable: Boolean = false
 
   override def dataType: DataType = LongType
@@ -88,6 +94,4 @@ case class MonotonicallyIncreasingID() extends LeafExpression with Stateful {
   override def nodeName: String = "monotonically_increasing_id"
 
   override def sql: String = s"$prettyName()"
-
-  override def freshCopy(): MonotonicallyIncreasingID = MonotonicallyIncreasingID()
 }

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

@@ -33,16 +33,20 @@ class InterpretedProjection(expressions: Seq[Expression]) extends Projection {
   def this(expressions: Seq[Expression], inputSchema: Seq[Attribute]) =
     this(bindReferences(expressions, inputSchema))
 
+  // null check is required for when Kryo invokes the no-arg constructor.
+  protected val exprArray = if (expressions != null) {
+    prepareExpressions(expressions, subExprEliminationEnabled = false).toArray
+  } else {
+    null
+  }
+
   override def initialize(partitionIndex: Int): Unit = {
-    expressions.foreach(_.foreach {
+    exprArray.foreach(_.foreach {
       case n: Nondeterministic => n.initialize(partitionIndex)
       case _ =>
     })
   }
 
-  // null check is required for when Kryo invokes the no-arg constructor.
-  protected val exprArray = if (expressions != null) expressions.toArray else null
-
   def apply(input: InternalRow): InternalRow = {
     val outputArray = new Array[Any](exprArray.length)
     var i = 0

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

@@ -58,6 +58,11 @@ case class ScalaUDF(
 
   override lazy val deterministic: Boolean = udfDeterministic && children.forall(_.deterministic)
 
+  // `ScalaUDF` uses `ExpressionEncoder` to convert the function result to Catalyst internal format.
+  // `ExpressionEncoder` is stateful as it reuses the `UnsafeRow` instance, thus `ScalaUDF` is
+  // stateful as well.
+  override def stateful: Boolean = true
+
   final override val nodePatterns: Seq[TreePattern] = Seq(SCALA_UDF)
 
   override def toString: String = s"$name(${children.mkString(", ")})"

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala

@@ -1271,8 +1271,11 @@ class CodegenContext extends Logging {
   def generateExpressions(
       expressions: Seq[Expression],
       doSubexpressionElimination: Boolean = false): Seq[ExprCode] = {
-    if (doSubexpressionElimination) subexpressionElimination(expressions)
-    expressions.map(e => e.genCode(this))
+    // We need to make sure that we do not reuse stateful expressions. This is needed for codegen
+    // as well because some expressions may implement `CodegenFallback`.
+    val cleanedExpressions = expressions.map(_.freshCopyIfContainsStatefulExpression())
+    if (doSubexpressionElimination) subexpressionElimination(cleanedExpressions)
+    cleanedExpressions.map(e => e.genCode(this))
   }
 
   /**

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

@@ -1125,11 +1125,13 @@ case class SortArray(base: Expression, ascendingOrder: Expression)
   """,
   group = "array_funcs",
   since = "2.4.0")
-case class Shuffle(child: Expression, randomSeed: Option[Long] = None)
-  extends UnaryExpression with ExpectsInputTypes with Stateful with ExpressionWithRandomSeed {
+case class Shuffle(child: Expression, randomSeed: Option[Long] = None) extends UnaryExpression
+  with ExpectsInputTypes with Nondeterministic with ExpressionWithRandomSeed {
 
   def this(child: Expression) = this(child, None)
 
+  override def stateful: Boolean = true
+
   override def seedExpression: Expression = randomSeed.map(Literal.apply).getOrElse(UnresolvedSeed)
 
   override def withNewSeed(seed: Long): Shuffle = copy(randomSeed = Some(seed))
@@ -1195,8 +1197,6 @@ case class Shuffle(child: Expression, randomSeed: Option[Long] = None)
      """.stripMargin
   }
 
-  override def freshCopy(): Shuffle = Shuffle(child, randomSeed)
-
   override def withNewChildInternal(newChild: Expression): Shuffle = copy(child = newChild)
 }
 

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

@@ -201,7 +201,7 @@ case class CurrentCatalog() extends LeafExpression with Unevaluable {
   since = "2.3.0",
   group = "misc_funcs")
 // scalastyle:on line.size.limit
-case class Uuid(randomSeed: Option[Long] = None) extends LeafExpression with Stateful
+case class Uuid(randomSeed: Option[Long] = None) extends LeafExpression with Nondeterministic
     with ExpressionWithRandomSeed {
 
   def this() = this(None)
@@ -216,6 +216,8 @@ case class Uuid(randomSeed: Option[Long] = None) extends LeafExpression with Sta
 
   override def dataType: DataType = StringType
 
+  override def stateful: Boolean = true
+
   @transient private[this] var randomGenerator: RandomUUIDGenerator = _
 
   override protected def initializeInternal(partitionIndex: Int): Unit =
@@ -235,8 +237,6 @@ case class Uuid(randomSeed: Option[Long] = None) extends LeafExpression with Sta
     ev.copy(code = code"final UTF8String ${ev.value} = $randomGen.getNextUUIDUTF8String();",
       isNull = FalseLiteral)
   }
-
-  override def freshCopy(): Uuid = Uuid(randomSeed)
 }
 
 // scalastyle:off line.size.limit