如何保留Spark HashingTF()函数的输入键或索引?

时间:2015-07-01 01:28:13

标签: apache-spark apache-spark-mllib tf-idf

基于1.4(https://spark.apache.org/docs/1.4.0/mllib-feature-extraction.html)的Spark文档,我正在编写一个TF-IDF示例,用于将文本文档转换为值向量。给出的示例显示了如何完成此操作,但输入是令牌的RDD,没有键。这意味着我的输出RDD不再包含索引或键以引用回原始文档。例如:

documents = sc.textFile("...").map(lambda line: line.split(" "))

hashingTF = HashingTF()
tf = hashingTF.transform(documents)

我想做这样的事情:

documents = sc.textFile("...").map(lambda line: (UNIQUE_LINE_KEY, line.split(" ")))

hashingTF = HashingTF()
tf = hashingTF.transform(documents)

并让生成的tf变量在某处包含UNIQUE_LINE_KEY值。我只是缺少一些明显的东西吗?从示例中可以看出,没有好方法将document RDD与tf RDD相关联。

2 个答案:

答案 0 :(得分:3)

我也遇到了同样的问题。在文档的示例中,他们鼓励您直接在RDD上应用转换。

但是,您可以在矢量本身上应用变换,这样您就可以按照您选择的方式保留键。

val input = sc.textFile("...")
val documents = input.map(doc => doc -> doc.split(" ").toSeq)

val hashingTF = new HashingTF()
val tf = documents.mapValues(hashingTF.transform(_))
tf.cache()
val idf = new IDF().fit(tf.values)
val tfidf = tf.mapValues(idf.transform(_))

请注意,此代码将生成RDD [(String,Vector)]而不是RDD [Vector]

答案 1 :(得分:2)

如果您在提交85b96372cf0fd055f89fc639f45c1f2cb02a378f之后使用Spark版本(这包括1.4),并使用ml API HashingTF(需要DataFrame输入而不是普通RDD),则输出中的原始列。希望有所帮助!