如何在Spark中的每一行添加源文件名？

Question

我是Spark的新手，我正在尝试使用它来自的文件名向每个输入行插入一列。

我见过其他人问了一个类似的问题，但是他们所有的答案都使用了wholeTextFile，但我正在尝试为更大的CSV文件（使用Spark-CSV库读取），JSON文件和Parquet文件（不仅仅是小文本）执行此操作文件）。

我可以使用spark-shell获取文件名列表：

val df = sqlContext.read.parquet("/blah/dir")
val names = df.select(inputFileName())
names.show

但这是一个数据帧。我不确定如何将它作为列添加到每一行（如果该结果的排序与初始数据相同，但我认为它总是如此）以及如何将其作为所有输入类型的通用解决方案。

Answer 1

我刚刚发现的另一个解决方案是将文件名添加为DataFrame中的一列

val df = sqlContext.read.parquet("/blah/dir")

val dfWithCol = df.withColumn("filename",input_file_name())

参考：spark load data and add filename as dataframe column

Answer 2

从文本文件创建RDD时，您可能希望将数据映射到案例类，因此可以在该阶段添加输入源：

case class Person(inputPath: String, name: String, age: Int)
val inputPath = "hdfs://localhost:9000/tmp/demo-input-data/persons.txt"
val rdd = sc.textFile(inputPath).map {
    l =>
      val tokens = l.split(",")
      Person(inputPath, tokens(0), tokens(1).trim().toInt)
  }
rdd.collect().foreach(println)

如果您不想将“业务数据”与元数据混合：

case class InputSourceMetaData(path: String, size: Long)
case class PersonWithMd(name: String, age: Int, metaData: InputSourceMetaData)

// Fake the size, for demo purposes only
val md = InputSourceMetaData(inputPath, size = -1L)
val rdd = sc.textFile(inputPath).map {
  l =>
    val tokens = l.split(",")
    PersonWithMd(tokens(0), tokens(1).trim().toInt, md)
}
rdd.collect().foreach(println)

如果您将RDD提升为DataFrame：

import sqlContext.implicits._
val df = rdd.toDF()
df.registerTempTable("x")

你可以查询它

sqlContext.sql("select name, metadata from x").show()
sqlContext.sql("select name, metadata.path from x").show()
sqlContext.sql("select name, metadata.path, metadata.size from x").show()

更新

您可以使用org.apache.hadoop.fs.FileSystem.listFiles()递归读取HDFS中的文件。

给定值files（包含org.apache.hadoop.fs.LocatedFileStatus的标准Scala集合）中的文件名列表，您可以为每个文件创建一个RDD：

val rdds = files.map { f =>
  val md = InputSourceMetaData(f.getPath.toString, f.getLen)

  sc.textFile(md.path).map {
    l =>
      val tokens = l.split(",")
      PersonWithMd(tokens(0), tokens(1).trim().toInt, md)
  }
}

现在你可以将RDD的列表reduce合并为一个：reduce的函数将所有RDD串联成一个：

val rdd = rdds.reduce(_ ++ _)
rdd.collect().foreach(println)

这有效，但我无法测试这是否与大文件分配/执行良好。