Add comment and move tests

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

@@ -382,6 +382,7 @@ class CometSparkSessionExtensions
             case None => b
           }
 
+        // For AQE shuffle stage on a Comet shuffle exchange
         case s @ ShuffleQueryStageExec(_, _: CometShuffleExchangeExec, _) =>
           val newOp = transform1(s)
           newOp match {
@@ -391,6 +392,9 @@ class CometSparkSessionExtensions
               s
           }
 
+        // For AQE shuffle stage on a reused Comet shuffle exchange
+        // Note that we don't need to handle `ReusedExchangeExec` for non-AQE case, because
+        // the query plan won't be re-optimized/planned in non-AQE mode.
         case s @ ShuffleQueryStageExec(
               _,
               ReusedExchangeExec(_, _: CometShuffleExchangeExec),

spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala

@@ -926,6 +926,12 @@ class CometAsyncShuffleSuite extends CometColumnarShuffleSuite {
   override protected val asyncShuffleEnable: Boolean = true
 
   protected val adaptiveExecutionEnabled: Boolean = true
+}
+
+class CometShuffleSuite extends CometColumnarShuffleSuite {
+  override protected val asyncShuffleEnable: Boolean = false
+
+  protected val adaptiveExecutionEnabled: Boolean = true
 
   import testImplicits._
 
@@ -976,12 +982,6 @@ class CometAsyncShuffleSuite extends CometColumnarShuffleSuite {
   }
 }
 
-class CometShuffleSuite extends CometColumnarShuffleSuite {
-  override protected val asyncShuffleEnable: Boolean = false
-
-  protected val adaptiveExecutionEnabled: Boolean = true
-}
-
 class DisableAQECometShuffleSuite extends CometColumnarShuffleSuite {
   override protected val asyncShuffleEnable: Boolean = false