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..44f383cf0188 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 @@ -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)) { 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 802e86374f50..c5e90a8ba852 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 @@ -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) } 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..5b6dfc71e362 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,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. */ 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..3c0cd19f7793 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,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) + } } 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] {