Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
marcuslin123 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)) =>
Expand All @@ -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)
}
Expand Down
18 changes: 15 additions & 3 deletions .../src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToLocalRelationSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
}
}


Expand Down
10 changes: 8 additions & 2 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -308,15 +308,21 @@ class QueryExecution(

def assertCommandExecuted(): Unit = commandExecuted

private def cloneWithFreshStatefulExpressions(plan: LogicalPlan): LogicalPlan = {
plan.clone().transformWithSubqueries {
case node => node.mapExpressions(_.freshCopyIfContainsStatefulExpression())

@vinodkc vinodkc Aug 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

mapExpressions uses fastEquals to decide whether an expression actually changed:

val newE = f(e)
if (newE.fastEquals(e)) e        // structural equality — not reference equality
else { changed = true; newE }

This works for NamedLambdaVariable because value: AtomicReference is a constructor argument — two NLVs with different AtomicReference instances are not structurally equal, so fastEquals returns false and the fresh copy sticks.

But for expressions like RegExpReplace, StringTranslate, FormatNumber (SPARK-58204), the mutable state lives in a @transient var — not in the constructor. A fresh copy of RegExpReplace has identical constructor args (subject, regexp, replacement), so fastEquals returns true and the fresh copy is silently discarded — the deep-copy becomes a no-op for those expressions.

In practice this is not a correctness issue today — ConvertToLocalRelation has its own direct freshCopyIfContainsStatefulExpression() call that covers driver-side evaluation, and non-local plans serialize fresh copies to each executor task. But it is a latent gap for future stateful expressions.

A simple fix is to use reference equality (ne) instead:

// change detection with ne instead of fastEquals
val newE = e.freshCopyIfContainsStatefulExpression()
if (newE ne e) { changed = true; newE } else e

}
}

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
Expand Down
25 changes: 24 additions & 1 deletion sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, 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
Expand Down Expand Up @@ -55,6 +55,14 @@ 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
}

def checkDumpedPlans(path: String, expected: Int): Unit = Utils.tryWithResource(
Source.fromFile(path)) { source =>
assert(source.getLines().toList
Expand Down Expand Up @@ -105,6 +113,21 @@ 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("dumping query execution info by invalid path") {
val path = "1234567890://plans.txt"
val exception = intercept[IllegalArgumentException] {
Expand Down