From 23db0925b1f7c08a8f37854e1dee140bb7c7c17d Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sun, 30 Aug 2026 20:41:35 -0700 Subject: [PATCH 1/2] [AURON #2491] Carry Spark's logical link onto converted native plans Converting a plan node builds a new node, and Spark's logical link does not travel with it. AQE asserts that a new query stage's subtree carries one. Spark 3 never reaches that assertion on a repeated collect, because the adaptive plan returns its finalized form unchanged; Spark 4 rebuilds query stages on every collect and sees the plan after the columnar rules have run, so a fully native plan with no exchange in it carries no link anywhere and the assertion fires. Set the link on the node that replaces the original, alongside the tags the recursive walk already carries across. Only a replacement needs it: a node conversion returned unchanged already holds its link, and setting it again would promote an inherited tag to a primary one. The walk is the right place rather than tryConvert, which is not the only path that builds a replacement and which also runs during the throwaway conversion that computes tags. --- .../org/apache/auron/AuronQuerySuite.scala | 75 +++++++++++++++++++ .../spark/sql/auron/AuronConverters.scala | 9 ++- 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala index 863f9c26a..5240991a4 100644 --- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala +++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala @@ -18,7 +18,10 @@ package org.apache.auron import org.apache.spark.sql.{AuronQueryTest, Row} import org.apache.spark.sql.auron.NativeRDD +import org.apache.spark.sql.auron.NativeSupports import org.apache.spark.sql.auron.join.JoinBuildSides.{JoinBuildLeft, JoinBuildRight} +import org.apache.spark.sql.catalyst.plans.logical.{Join, LogicalPlan} +import org.apache.spark.sql.execution.auron.plan.NativeBroadcastJoinBase import org.apache.spark.sql.execution.auron.plan.NativeFilterBase import org.apache.spark.sql.execution.auron.plan.NativeShuffledHashJoinBase import org.apache.spark.sql.execution.auron.plan.NativeShuffleExchangeBase @@ -1144,4 +1147,76 @@ class AuronQuerySuite extends AuronQueryTest with BaseAuronSQLSuite with AuronSQ checkSparkAnswerAndOperator("select * from orc_or where id = 5 or b = 900000") } } + + private def withLinkTestTables(body: => Unit): Unit = { + withTable("t1", "t2") { + sql("create table t1 using parquet as select id as c1, id + 1 as c2 from range(10)") + sql("create table t2 using parquet as select id as c3 from range(5)") + body + } + } + + test("the same DataFrame can be executed more than once") { + withLinkTestTables { + val df = + checkSparkAnswerAndOperator("select c1 from t1 where c2 > (select max(c3) from t2)") + + // Re-executing the same DataFrame re-enters adaptive stage creation, this time against + // the plan Auron produced rather than the plan Spark planned. + checkAnswer(df, (4L to 9L).map(Row(_))) + } + } + + test("converted native plans carry Spark's logical link") { + withLinkTestTables { + val df = + checkSparkAnswerAndOperator("select c1 from t1 where c2 > (select max(c3) from t2)") + + val plan = stripAQEPlan(df.queryExecution.executedPlan) + val nativeNodes = plan.collect { case p: NativeSupports => p } + assert(nativeNodes.nonEmpty, s"expected native operators in:\n$plan") + + // Every native node is checked because this query needs no exchange in the main tree. + assert( + nativeNodes.forall(_.logicalLink.isDefined), + "native operators without a logical link: " + + s"${nativeNodes.filter(_.logicalLink.isEmpty).map(_.nodeName).mkString(", ")}\n$plan") + } + } + + test("converted native plans keep one logical link per operator") { + withLinkTestTables { + val df = + checkSparkAnswerAndOperator("select t1.c1, t2.c3 from t1 join t2 on t1.c2 = t2.c3") + + val plan = stripAQEPlan(df.queryExecution.executedPlan) + val links = plan.collect { case p: NativeSupports => p }.flatMap(_.logicalLink) + val distinctLinks = links.foldLeft(Seq.empty[LogicalPlan]) { (acc, link) => + if (acc.exists(_.eq(link))) acc else acc :+ link + } + + // A strategy plans one logical node into a whole physical subtree and links only that + // subtree's root, so nodes within one planning unit legitimately share a link. Across + // units the links must differ. Comparing by reference keeps two units of the same kind + // distinct. + assert( + distinctLinks.size > 1, + "native operators share a single logical link: " + + s"${distinctLinks.map(_.nodeName).mkString(", ")}\n$plan") + + // Counting links cannot tell a correct labelling from one shifted a level up the tree, + // where every node carries its parent's link and the count is unchanged. Pin the join. + val joinNodes = plan.collect { + case p: NativeBroadcastJoinBase => p + case p: NativeSortMergeJoinBase => p + case p: NativeShuffledHashJoinBase => p + } + assert(joinNodes.size == 1, s"expected exactly one native join in:\n$plan") + assert( + joinNodes.head.logicalLink.exists(_.isInstanceOf[Join]), + "native join is linked to " + + s"${joinNodes.head.logicalLink.map(_.nodeName).getOrElse("nothing")}, not a Join" + + s"\n$plan") + } + } } diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala index b9847ddbc..06e46ef70 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala @@ -202,7 +202,14 @@ object AuronConverters extends Logging { exec.getTagValue(joinSmallerSideTag).foreach(newExec.setTagValue(joinSmallerSideTag, _)) if (!isNeverConvert(newExec)) { - newExec = convertSparkPlan(newExec) + val convertedExec = convertSparkPlan(newExec) + newExec = if (convertedExec.eq(newExec)) { + convertedExec + } else { + // A converted node is built from scratch, so Spark's logical link does not travel + // with it. AQE searches a query stage's subtree for a link and asserts it finds one. + Shims.get.setLogicalLink(convertedExec, exec) + } } danglingConverted = newDanglingConverted :+ newExec } From 2cf9541d8b516ae7b80cceb01ff796ea5728a73f Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sun, 30 Aug 2026 23:19:25 -0700 Subject: [PATCH 2/2] [AURON #2491] Restrict the logical link assertions to Spark 3.2 and above Before Spark 3.2, TreeNode.withNewChildren is overridable and reaches copyTagsFrom only by way of makeCopy. The native plan nodes override it with a plain copy of the case class, so no node tag survives a child rebuild on those versions. The converter sets the logical link on every version, but on 3.0 and 3.1 adaptive execution drops it again as soon as it substitutes a query stage into the tree, so the two assertions on the executed plan cannot hold there. They are now assumed away on those versions rather than asserted. The repeated execution test is unchanged and still runs everywhere. --- .../scala/org/apache/auron/AuronQuerySuite.scala | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala index 5240991a4..0f4f224c6 100644 --- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala +++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala @@ -1156,6 +1156,16 @@ class AuronQuerySuite extends AuronQueryTest with BaseAuronSQLSuite with AuronSQ } } + // Before Spark 3.2, TreeNode.withNewChildren is overridable and copies tags only by way of + // makeCopy. The native plan nodes override it with a plain copy of the case class, so no tag + // survives a child rebuild on those versions. The converter still sets the logical link + // everywhere, but adaptive execution drops it again as soon as it substitutes a query stage + // into the tree, so only the versions below can be asserted on the executed plan. + private def assumeTagsSurviveChildRebuild(): Unit = + assume( + sparkver.matchVersion("3.2 / 3.3 / 3.4 / 3.5 / 4.0 / 4.1 / 4.2"), + "native plan nodes drop tags on withNewChildren before Spark 3.2") + test("the same DataFrame can be executed more than once") { withLinkTestTables { val df = @@ -1168,6 +1178,7 @@ class AuronQuerySuite extends AuronQueryTest with BaseAuronSQLSuite with AuronSQ } test("converted native plans carry Spark's logical link") { + assumeTagsSurviveChildRebuild() withLinkTestTables { val df = checkSparkAnswerAndOperator("select c1 from t1 where c2 > (select max(c3) from t2)") @@ -1185,6 +1196,7 @@ class AuronQuerySuite extends AuronQueryTest with BaseAuronSQLSuite with AuronSQ } test("converted native plans keep one logical link per operator") { + assumeTagsSurviveChildRebuild() withLinkTestTables { val df = checkSparkAnswerAndOperator("select t1.c1, t2.c3 from t1 join t2 on t1.c2 = t2.c3")