AdaptiveQueryExecSuite.scala

class AdaptiveQueryExecSuite
  extends QueryTest
  with SharedSparkSession
  with AdaptiveSparkPlanHelper {
  
  test("Empty stage") {
    withSQLConf(
      SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
      SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true") {
      withTempView("v1", "v2") {
        spark.range(10).withColumn("a", 'id).createTempView("v1")
        spark.range(10).withColumn("b", 'id).createTempView("v2")
        withTable("t1", "t2") {
          sql("CREATE TABLE t1 USING PARQUET SELECT * FROM v1")
          sql("CREATE TABLE t2 USING PARQUET SELECT * FROM v2")
          val testDf = sql(
            """
              | SELECT b, COUNT(t1.a) as cnt
              | FROM t1
              | INNER JOIN t2
              | ON t1.id = t2.id
              | WHERE t1.id > 10
              | GROUP BY b
              |""".stripMargin)
          checkAnswer(testDf, Seq())
          val plan = testDf.queryExecution.executedPlan
          val coalescedReaders = collect(plan) {
            case r: CustomShuffleReaderExec => r
          }
          assert(coalescedReaders.length == 1, s"$plan")
          assert(coalescedReaders.head.outputPartitioning.numPartitions == 1)
        }
      }
    }
  }
}