Fix tests.

Author: nongli

sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegen.scala

@@ -39,6 +39,7 @@ trait CodegenSupport extends SparkPlan {
   private def variablePrefix: String = this match {
     case _: TungstenAggregate => "agg"
     case _: BroadcastHashJoin => "bhj"
+    case _: PhysicalRDD => "rddScan"
     case _ => nodeName.toLowerCase
   }
 

sql/hive/src/test/scala/org/apache/spark/sql/sources/BucketedReadSuite.scala

@@ -73,32 +73,34 @@ class BucketedReadSuite extends QueryTest with SQLTestUtils with TestHiveSinglet
       bucketValues: Seq[Integer],
       filterCondition: Column,
       originalDataFrame: DataFrame): Unit = {
+    // This test verifies parts of the plan. Disable whole stage codegen.
+    withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false") {
+      val bucketedDataFrame = hiveContext.table("bucketed_table").select("i", "j", "k")
+      val BucketSpec(numBuckets, bucketColumnNames, _) = bucketSpec
+      // Limit: bucket pruning only works when the bucket column has one and only one column
+      assert(bucketColumnNames.length == 1)
+      val bucketColumnIndex = bucketedDataFrame.schema.fieldIndex(bucketColumnNames.head)
+      val bucketColumn = bucketedDataFrame.schema.toAttributes(bucketColumnIndex)
+      val matchedBuckets = new BitSet(numBuckets)
+      bucketValues.foreach { value =>
+        matchedBuckets.set(DataSourceStrategy.getBucketId(bucketColumn, numBuckets, value))
+      }
 
-    val bucketedDataFrame = hiveContext.table("bucketed_table").select("i", "j", "k")
-    val BucketSpec(numBuckets, bucketColumnNames, _) = bucketSpec
-    // Limit: bucket pruning only works when the bucket column has one and only one column
-    assert(bucketColumnNames.length == 1)
-    val bucketColumnIndex = bucketedDataFrame.schema.fieldIndex(bucketColumnNames.head)
-    val bucketColumn = bucketedDataFrame.schema.toAttributes(bucketColumnIndex)
-    val matchedBuckets = new BitSet(numBuckets)
-    bucketValues.foreach { value =>
-      matchedBuckets.set(DataSourceStrategy.getBucketId(bucketColumn, numBuckets, value))
-    }
+      // Filter could hide the bug in bucket pruning. Thus, skipping all the filters
+      val plan = bucketedDataFrame.filter(filterCondition).queryExecution.executedPlan
+      val rdd = plan.find(_.isInstanceOf[PhysicalRDD])
+      assert(rdd.isDefined, plan)
 
-    // Filter could hide the bug in bucket pruning. Thus, skipping all the filters
-    val rdd = bucketedDataFrame.filter(filterCondition).queryExecution.executedPlan
-      .find(_.isInstanceOf[PhysicalRDD])
-    assert(rdd.isDefined)
+      val checkedResult = rdd.get.execute().mapPartitionsWithIndex { case (index, iter) =>
+        if (matchedBuckets.get(index % numBuckets)) Iterator(true) else Iterator(iter.isEmpty)
+      }
+      // checking if all the pruned buckets are empty
+      assert(checkedResult.collect().forall(_ == true))
 
-    val checkedResult = rdd.get.execute().mapPartitionsWithIndex { case (index, iter) =>
-      if (matchedBuckets.get(index % numBuckets)) Iterator(true) else Iterator(iter.isEmpty)
+      checkAnswer(
+        bucketedDataFrame.filter(filterCondition).orderBy("i", "j", "k"),
+        originalDataFrame.filter(filterCondition).orderBy("i", "j", "k"))
     }
-    // checking if all the pruned buckets are empty
-    assert(checkedResult.collect().forall(_ == true))
-
-    checkAnswer(
-      bucketedDataFrame.filter(filterCondition).orderBy("i", "j", "k"),
-      originalDataFrame.filter(filterCondition).orderBy("i", "j", "k"))
   }
 
   test("read partitioning bucketed tables with bucket pruning filters") {