Parallel list files on S3 with Spark

gistfile1.scala

import org.apache.hadoop.fs.{FileSystem, Path}
import org.apache.hadoop.conf.Configuration 

  case class S3File(path: String, isDir: Boolean, size: Long) {
    def children = listFiles(path)
  }

  def listFiles(path: String): Seq[S3File] = {
    val fs = FileSystem.get(new java.net.URI(path), new Configuration())
    fs.listStatus(new Path(path)).map(s => S3File(s.getPath.toString, s.isDir, s.getLen))
  }

  def recursivelyListFiles(path: String): Seq[S3File] = recursivelyListFiles(path :: Nil)

  def recursivelyListFiles(paths: Seq[String]): Seq[S3File] = {
    val remainingDirectories = new scala.collection.mutable.ArrayBuffer[S3File]
    val allFiles = new scala.collection.mutable.ArrayBuffer[S3File]

    remainingDirectories ++= paths.map(S3File(_, isDir = true, 0))
    while (remainingDirectories.nonEmpty) {
      val newDirs = sparkContext.parallelize(remainingDirectories.map(_.path))
      val currentBatch = newDirs.mapPartitions { iter =>
        val fs = FileSystem.get(new java.net.URI(paths.head), new Configuration())
        iter.flatMap{path =>
          try {
            fs.listStatus(new Path(path)).map(s => S3File(s.getPath.toString, s.isDir, s.getLen))
          } catch {
            case e: java.io.FileNotFoundException =>
              println(s"File $path not found.")
              Nil
          }
        }
      }.collect()
      val (dirs, files) = currentBatch.partition(_.isDir)
      remainingDirectories.clear()
      remainingDirectories ++= dirs
      allFiles ++= files
    }
    allFiles
  }