This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 73d3dbcf71ac [SPARK-58208][SQL] Deep-copy stateful expressions before
optimization
73d3dbcf71ac is described below
commit 73d3dbcf71acbb6fecb15b4480462543eb1bff20
Author: Marcus Lin <[email protected]>
AuthorDate: Sat Aug 22 02:40:32 2026 +0800
[SPARK-58208][SQL] Deep-copy stateful expressions before optimization
### What changes were proposed in this pull request?
This PR deep-copies stateful expressions before they are evaluated by the
optimizer.
Specifically, it:
- clones the logical plan and replaces stateful expressions with fresh
copies before `QueryExecution.optimizedPlan` invokes the optimizer;
- uses fresh stateful expressions when `ConvertToLocalRelation` evaluates
project and filter expressions; and
- adds regression coverage for both paths using `NamedLambdaVariable`.
### Why are the changes needed?
Logical plans can be shared by DataFrames derived from the same base
DataFrame. Although `QueryExecution` cloned the plan before optimization,
stateful expressions such as `NamedLambdaVariable` could retain shared mutable
state across those clones. Concurrent optimization could therefore mutate the
same expression state from multiple threads and silently corrupt query results.
`ConvertToLocalRelation` had the same issue because it evaluated
expressions from the shared plan directly.
### Does this PR introduce _any_ user-facing change?
Yes. Concurrent actions on DataFrames derived from a shared logical plan no
longer risk silent result corruption caused by shared mutable expression state.
### How was this patch tested?
Added regression tests to `ConvertToLocalRelationSuite` and
`QueryExecutionSuite`.
```
build/sbt 'catalyst/testOnly
org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelationSuite'
build/sbt 'sql/testOnly org.apache.spark.sql.execution.QueryExecutionSuite'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5). Codex was used for code assistance and review;
the patch was not entirely generated by Codex.
Closes #57980 from
marcuslin123/codex/spark-58208-stateful-expression-copies.
Authored-by: Marcus Lin <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../catalyst/expressions/randomExpressions.scala | 6 +--
.../spark/sql/catalyst/optimizer/Optimizer.scala | 9 +++--
.../spark/sql/catalyst/plans/QueryPlan.scala | 42 +++++++++++++++++++-
.../apache/spark/sql/catalyst/trees/TreeNode.scala | 32 +++++++++++++++
.../sql/catalyst/expressions/RandomSuite.scala | 18 +++++++++
.../optimizer/ConvertToLocalRelationSuite.scala | 18 +++++++--
.../spark/sql/execution/QueryExecution.scala | 11 +++++-
.../spark/sql/execution/QueryExecutionSuite.scala | 46 +++++++++++++++++++++-
8 files changed, 169 insertions(+), 13 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala
index b52d09fc9c70..6a79ac339f1e 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/randomExpressions.scala
@@ -301,14 +301,14 @@ case class Uniform(
override def third: Expression = seedExpression
override def withNewSeed(newSeed: Long): Expression =
- Uniform(min, max, Literal(newSeed, LongType), hideSeed)
+ Uniform(min, max, Literal(newSeed, LongType), hideSeed, timeZoneId)
override def withShiftedSeed(shift: Long): Expression =
- Uniform(min, max, Literal(seed + shift, LongType), hideSeed)
+ Uniform(min, max, Literal(seed + shift, LongType), hideSeed, timeZoneId)
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)) {
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index 77ffd7e29da0..58b6479ddee3 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -2836,9 +2836,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)) =>
@@ -2849,7 +2851,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)
}
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
index 785131875114..ed4e889aa5d9 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala
@@ -258,13 +258,29 @@ 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)) {
+ // Reference equality preserves fresh stateful copies that fastEquals
sees as unchanged.
+ val unchanged = if (useFastEquals) newE.fastEquals(e) else newE.eq(e)
+ if (unchanged) {
e
} else {
changed = true
@@ -577,6 +593,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.
*/
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
index e82e6a30b9bb..e88b0c214c76 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreeNode.scala
+++
b/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
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala
index 9e6b59b51138..858c6e3b5028 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/RandomSuite.scala
@@ -64,4 +64,22 @@ 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)
+
+ Seq(uniform.withNewSeed(1), uniform.withShiftedSeed(1)).foreach {
+ case copied: Uniform =>
+ assert(copied.timeZoneId == uniform.timeZoneId)
+ assert(copied.resolved)
+ case other => fail(s"Expected Uniform but got ${other.getClass.getName}")
+ }
+ }
}
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToLocalRelationSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToLocalRelationSuite.scala
index 622af60d85d9..f4d412153b0e 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToLocalRelationSuite.scala
+++
b/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)
+ }
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
index 1fa70351f2d0..51e54ef5f753 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
@@ -308,6 +308,13 @@ 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
@@ -315,8 +322,8 @@ class QueryExecution(
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
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
index f7afdb5e6e53..fd047d54b957 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
+++
b/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] {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
