[GitHub] [spark] dongjoon-hyun commented on a diff in pull request #40589: [SPARK-38697][SQL] Extend SparkSessionExtensions to inject rules into AQE query stage optimizer

["dongjoon-hyun (via GitHub)" <gi...@apache.org>]

dongjoon-hyun commented on code in PR #40589:
URL: https://github.com/apache/spark/pull/40589#discussion_r1153984929


##########
sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala:
##########
@@ -500,6 +500,22 @@ class SparkSessionExtensionSuite extends SparkFunSuite with SQLHelper {
       }
     }
   }
+
+  test("SPARK-38697: Extend SparkSessionExtensions to inject rules into AQE query stage " +
+    "optimizer") {
+    val extensions = create { extensions =>
+      extensions.injectQueryStageOptimizerRule(_ => FailWithAQEShuffleRead)
+    }
+    withSession(extensions) { session =>
+      assert(session.sessionState.adaptiveRulesHolder.queryStageOptimizerRules
+        .contains(FailWithAQEShuffleRead))
+      val df = session.range(1).repartition()
+      intercept[Exception] {

Review Comment:
   This looks too general. Please check the exception message to make it sure.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org