diff --git a/gluten-core/src/main/scala/io/glutenproject/extension/ColumnarOverrides.scala b/gluten-core/src/main/scala/io/glutenproject/extension/ColumnarOverrides.scala index 3c12d0c56e22..f6c2b2ac90a9 100644 --- a/gluten-core/src/main/scala/io/glutenproject/extension/ColumnarOverrides.scala +++ b/gluten-core/src/main/scala/io/glutenproject/extension/ColumnarOverrides.scala @@ -45,6 +45,8 @@ import org.apache.spark.sql.execution.window.WindowExec import org.apache.spark.sql.hive.HiveTableScanExecTransformer import org.apache.spark.util.SparkRuleUtil +import scala.collection.mutable.ListBuffer + // This rule will conduct the conversion from Spark plan to the plan transformer. case class TransformPreOverrides(isAdaptiveContext: Boolean) extends Rule[SparkPlan] @@ -749,7 +751,10 @@ case class ColumnarOverrideRules(session: SparkSession) private val aqeStackTraceIndex = 14 // Holds the original plan for possible entire fallback. - private val localOriginalPlan = new ThreadLocal[SparkPlan] + private val localOriginalPlans: ThreadLocal[ListBuffer[SparkPlan]] = + ThreadLocal.withInitial(() => ListBuffer.empty[SparkPlan]) + private val localIsAdaptiveContextFlags: ThreadLocal[ListBuffer[Boolean]] = + ThreadLocal.withInitial(() => ListBuffer.empty[Boolean]) // Do not create rules in class initialization as we should access SQLConf // while creating the rules. At this time SQLConf may not be there yet. @@ -760,7 +765,10 @@ case class ColumnarOverrideRules(session: SparkSession) } def isAdaptiveContext: Boolean = - session.sparkContext.getLocalProperty(GLUTEN_IS_ADAPTIVE_CONTEXT).toBoolean + Option(session.sparkContext.getLocalProperty(GLUTEN_IS_ADAPTIVE_CONTEXT)) + .getOrElse("false") + .toBoolean || + localIsAdaptiveContextFlags.get().head private def setAdaptiveContext(): Unit = { val traceElements = Thread.currentThread.getStackTrace @@ -772,25 +780,27 @@ case class ColumnarOverrideRules(session: SparkSession) // columnar rule will be applied in adaptive execution context. This part of code // needs to be carefully checked when supporting higher versions of spark to make // sure the calling stack has not been changed. - session.sparkContext.setLocalProperty( - GLUTEN_IS_ADAPTIVE_CONTEXT, - traceElements(aqeStackTraceIndex).getClassName - .equals(AdaptiveSparkPlanExec.getClass.getName) - .toString) + localIsAdaptiveContextFlags + .get() + .prepend( + traceElements(aqeStackTraceIndex).getClassName + .equals(AdaptiveSparkPlanExec.getClass.getName)) } private def resetAdaptiveContext(): Unit = - session.sparkContext.setLocalProperty(GLUTEN_IS_ADAPTIVE_CONTEXT, null) + localIsAdaptiveContextFlags.get().remove(0) - private def setOriginalPlan(plan: SparkPlan): Unit = localOriginalPlan.set(plan) + private def setOriginalPlan(plan: SparkPlan): Unit = { + localOriginalPlans.get.prepend(plan) + } private def originalPlan: SparkPlan = { - val plan = localOriginalPlan.get() + val plan = localOriginalPlans.get.head assert(plan != null) plan } - private def resetOriginalPlan(): Unit = localOriginalPlan.remove() + private def resetOriginalPlan(): Unit = localOriginalPlans.get.remove(0) private def preOverrides(): List[SparkSession => Rule[SparkPlan]] = { val tagBeforeTransformHitsRules = diff --git a/gluten-delta/src/test/scala/io/glutenproject/execution/VeloxDeltaSuite.scala b/gluten-delta/src/test/scala/io/glutenproject/execution/VeloxDeltaSuite.scala index 3f1f7ab255fa..7e096151b16c 100644 --- a/gluten-delta/src/test/scala/io/glutenproject/execution/VeloxDeltaSuite.scala +++ b/gluten-delta/src/test/scala/io/glutenproject/execution/VeloxDeltaSuite.scala @@ -55,4 +55,22 @@ class VeloxDeltaSuite extends WholeStageTransformerSuite { checkLengthAndPlan(df2, 1) checkAnswer(df2, Row("v2") :: Nil) } + + test("basic test with stats.skipping disabled") { + withSQLConf("spark.databricks.delta.stats.skipping" -> "false") { + spark.sql(s""" + |create table delta_test2 (id int, name string) using delta + |""".stripMargin) + spark.sql(s""" + |insert into delta_test2 values (1, "v1"), (2, "v2") + |""".stripMargin) + val df1 = runQueryAndCompare("select * from delta_test2") { _ => } + checkLengthAndPlan(df1, 2) + checkAnswer(df1, Row(1, "v1") :: Row(2, "v2") :: Nil) + + val df2 = runQueryAndCompare("select name from delta_test2 where id = 2") { _ => } + checkLengthAndPlan(df2, 1) + checkAnswer(df2, Row("v2") :: Nil) + } + } }