add minSplits for WholeTextFiles

Author: yinxusen

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -456,12 +456,13 @@ class SparkContext(config: SparkConf) extends Logging {
    *
    * @note Small files are preferred, as each file will be loaded fully in memory.
    */
-  def wholeTextFiles(path: String): RDD[(String, String)] = {
+  def wholeTextFiles(path: String, minSplits: Int = defaultMinSplits): RDD[(String, String)] = {
     newAPIHadoopFile(
       path,
       classOf[WholeTextFileInputFormat],
       classOf[String],
-      classOf[String])
+      classOf[String],
+      minSplits = minSplits)
   }
 
   /**
@@ -584,11 +585,12 @@ class SparkContext(config: SparkConf) extends Logging {
       fClass: Class[F],
       kClass: Class[K],
       vClass: Class[V],
-      conf: Configuration = hadoopConfiguration): RDD[(K, V)] = {
+      conf: Configuration = hadoopConfiguration,
+      minSplits: Int = 1): RDD[(K, V)] = {
     val job = new NewHadoopJob(conf)
     NewFileInputFormat.addInputPath(job, new Path(path))
     val updatedConf = job.getConfiguration
-    new NewHadoopRDD(this, fClass, kClass, vClass, updatedConf)
+    new NewHadoopRDD(this, fClass, kClass, vClass, updatedConf, minSplits)
   }
 
   /**

core/src/main/scala/org/apache/spark/input/WholeTextFileInputFormat.scala

@@ -17,7 +17,10 @@
 
 package org.apache.spark.input
 
+import scala.collection.JavaConversions._
+
 import org.apache.hadoop.fs.Path
+import org.apache.hadoop.fs.FileStatus
 import org.apache.hadoop.mapreduce.InputSplit
 import org.apache.hadoop.mapreduce.JobContext
 import org.apache.hadoop.mapreduce.lib.input.CombineFileInputFormat
@@ -44,4 +47,15 @@ private[spark] class WholeTextFileInputFormat extends CombineFileInputFormat[Str
       context,
       classOf[WholeTextFileRecordReader])
   }
+
+  /**
+   * Allow minSplits set by end-user in order to keep compatibility with old Hadoop API.
+   */
+  def setMaxSplitSize(context: JobContext, minSplits: Int) {
+    val files = listStatus(context)
+    val totalLen = files.map { file =>
+      if (file.isDir) 0L else file.getLen
+    }.sum
+    super.setMaxSplitSize(totalLen / (if (minSplits == 0) 1 else minSplits))
+  }
 }

core/src/main/scala/org/apache/spark/rdd/NewHadoopRDD.scala

@@ -26,6 +26,7 @@ import org.apache.hadoop.mapreduce._
 
 import org.apache.spark.{InterruptibleIterator, Logging, Partition, SerializableWritable, SparkContext, TaskContext}
 import org.apache.spark.annotation.DeveloperApi
+import org.apache.spark.input.WholeTextFileInputFormat
 
 private[spark]
 class NewHadoopPartition(rddId: Int, val index: Int, @transient rawSplit: InputSplit with Writable)
@@ -56,7 +57,8 @@ class NewHadoopRDD[K, V](
     inputFormatClass: Class[_ <: InputFormat[K, V]],
     keyClass: Class[K],
     valueClass: Class[V],
-    @transient conf: Configuration)
+    @transient conf: Configuration,
+    minSplits: Int = 1)
   extends RDD[(K, V)](sc, Nil)
   with SparkHadoopMapReduceUtil
   with Logging {
@@ -74,10 +76,17 @@ class NewHadoopRDD[K, V](
 
   override def getPartitions: Array[Partition] = {
     val inputFormat = inputFormatClass.newInstance
+
     if (inputFormat.isInstanceOf[Configurable]) {
       inputFormat.asInstanceOf[Configurable].setConf(conf)
     }
+
     val jobContext = newJobContext(conf, jobId)
+
+    if (inputFormat.isInstanceOf[WholeTextFileInputFormat]) {
+      inputFormat.asInstanceOf[WholeTextFileInputFormat].setMaxSplitSize(jobContext, minSplits)
+    }
+
     val rawSplits = inputFormat.getSplits(jobContext).toArray
     val result = new Array[Partition](rawSplits.size)
     for (i <- 0 until rawSplits.size) {