diff --git a/core/src/main/scala/org/apache/spark/rdd/PartitionPruningRDD.scala b/core/src/main/scala/org/apache/spark/rdd/PartitionPruningRDD.scala
index 2738a0089409a7c1bc39deef85ed4592f8df0d3f..574dd4233fb2724d62510307549196481611f111 100644
--- a/core/src/main/scala/org/apache/spark/rdd/PartitionPruningRDD.scala
+++ b/core/src/main/scala/org/apache/spark/rdd/PartitionPruningRDD.scala
@@ -33,8 +33,8 @@ class PruneDependency[T](rdd: RDD[T], @transient partitionFilterFunc: Int => Boo
   extends NarrowDependency[T](rdd) {
 
   @transient
-  val partitions: Array[Partition] = rdd.partitions.zipWithIndex
-    .filter(s => partitionFilterFunc(s._2)).map(_._1).zipWithIndex
+  val partitions: Array[Partition] = rdd.partitions
+    .filter(s => partitionFilterFunc(s.index)).zipWithIndex
     .map { case(split, idx) => new PartitionPruningRDDPartition(idx, split) : Partition }
 
   override def getParents(partitionId: Int) = {
diff --git a/core/src/test/scala/org/apache/spark/PartitionPruningRDDSuite.scala b/core/src/test/scala/org/apache/spark/rdd/PartitionPruningRDDSuite.scala
similarity index 92%
rename from core/src/test/scala/org/apache/spark/PartitionPruningRDDSuite.scala
rename to core/src/test/scala/org/apache/spark/rdd/PartitionPruningRDDSuite.scala
index 28e71e835fa2cf2bc92e8a8e4e4e600e392c8ea5..53a7b7c44df1c8345d0656c79f734cce6a1b47cd 100644
--- a/core/src/test/scala/org/apache/spark/PartitionPruningRDDSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rdd/PartitionPruningRDDSuite.scala
@@ -15,11 +15,10 @@
  * limitations under the License.
  */
 
-package org.apache.spark
+package org.apache.spark.rdd
 
 import org.scalatest.FunSuite
-import org.apache.spark.SparkContext._
-import org.apache.spark.rdd.{PartitionPruningRDDPartition, RDD, PartitionPruningRDD}
+import org.apache.spark.{TaskContext, Partition, SharedSparkContext}
 
 
 class PartitionPruningRDDSuite extends FunSuite with SharedSparkContext {
@@ -49,7 +48,7 @@ class PartitionPruningRDDSuite extends FunSuite with SharedSparkContext {
   }
 
 
-  test("Pruned Partitions can be merged ") {
+  test("Pruned Partitions can be unioned ") {
 
     val rdd = new RDD[Int](sc, Nil) {
       override protected def getPartitions = {
@@ -72,17 +71,11 @@ class PartitionPruningRDDSuite extends FunSuite with SharedSparkContext {
     })
 
     val merged = prunedRDD1 ++ prunedRDD2
-
     assert(merged.count() == 2)
     val take = merged.take(2)
-
     assert(take.apply(0) == 4)
-
     assert(take.apply(1) == 6)
-
-
   }
-
 }
 
 class TestPartition(i: Int, value: Int) extends Partition with Serializable {
@@ -90,4 +83,4 @@ class TestPartition(i: Int, value: Int) extends Partition with Serializable {
 
   def testValue = this.value
 
-}
\ No newline at end of file
+}