diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala index 1dfff412d9a8e..be90d3fb0361c 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala @@ -147,7 +147,7 @@ class EquivalentExpressions { expr.find(_.isInstanceOf[LambdaVariable]).isDefined || // `PlanExpression` wraps query plan. To compare query plans of `PlanExpression` on executor, // can cause error like NPE. - (expr.isInstanceOf[PlanExpression[_]] && TaskContext.get != null) + (expr.find(_.isInstanceOf[PlanExpression[_]]).isDefined && TaskContext.get != null) if (!skip && !addFunc(expr)) { childrenToRecurse(expr).foreach(addExprTree(_, addFunc)) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala index 65671d253dc53..b80843fee7881 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/SubexpressionEliminationSuite.scala @@ -16,8 +16,9 @@ */ package org.apache.spark.sql.catalyst.expressions -import org.apache.spark.SparkFunSuite +import org.apache.spark.{SparkFunSuite, TaskContext, TaskContextImpl} import org.apache.spark.sql.catalyst.expressions.codegen._ +import org.apache.spark.sql.catalyst.plans.logical.Command import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{BinaryType, DataType, IntegerType} @@ -255,6 +256,21 @@ class SubexpressionEliminationSuite extends SparkFunSuite with ExpressionEvalHel assert(equivalence2.getAllEquivalentExprs.count(_.size == 2) == 0) } + test("SPARK-38333: DPP expression should not be eliminated") { + try { + // support we in executor + val context1 = new TaskContextImpl(0, 0, 0, 0, 0, null, null, null) + TaskContext.setTaskContext(context1) + + val equivalence = new EquivalentExpressions + val expression = DynamicPruningExpression(Exists(TestCommand("foo"))) + equivalence.addExprTree(expression) + assert(equivalence.getEquivalentExprs(expression).size == 0) + } finally { + TaskContext.unset() + } + } + test("SPARK-34723: Correct parameter type for subexpression elimination under whole-stage") { withSQLConf(SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> "1") { val str = BoundReference(0, BinaryType, false) @@ -315,3 +331,5 @@ case class CodegenFallbackExpression(child: Expression) extends UnaryExpression with CodegenFallback { override def dataType: DataType = child.dataType } + +case class TestCommand(foo: String) extends Command