diff --git a/core/src/test/scala/org/apache/spark/api/python/PythonRDDSuite.scala b/core/src/test/scala/org/apache/spark/api/python/PythonRDDSuite.scala
index 1bebfe5ec84ec2b23bdb8d7931823ae7e727945f..5bcebabc9a4199fb31a77444790c8a383be19978 100644
--- a/core/src/test/scala/org/apache/spark/api/python/PythonRDDSuite.scala
+++ b/core/src/test/scala/org/apache/spark/api/python/PythonRDDSuite.scala
@@ -19,7 +19,6 @@ package org.apache.spark.api.python
 
 import org.scalatest.FunSuite
 import org.scalatest.matchers.ShouldMatchers
-import org.apache.spark.api.python.PythonRDD
 
 import java.io.{ByteArrayOutputStream, DataOutputStream}
 
diff --git a/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala b/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala
index cd01303bad0a08d9a19d8c972908f727e4f49ba1..223ebec5fa3dce3f739c915d0c8073b1dc4f64b4 100644
--- a/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala
+++ b/core/src/test/scala/org/apache/spark/rdd/RDDSuite.scala
@@ -48,7 +48,7 @@ class RDDSuite extends FunSuite with SharedSparkContext {
     val partitionSums = nums.mapPartitions(iter => Iterator(iter.reduceLeft(_ + _)))
     assert(partitionSums.collect().toList === List(3, 7))
 
-    val partitionSumsWithSplit = nums.mapPartitionsWithSplit {
+    val partitionSumsWithSplit = nums.mapPartitionsWithIndex {
       case(split, iter) => Iterator((split, iter.reduceLeft(_ + _)))
     }
     assert(partitionSumsWithSplit.collect().toList === List((0, 3), (1, 7)))