Spark 2.1.0 入门：特征抽取–CountVectorizer

【版权声明】博客内容由厦门大学数据库实验室拥有版权，未经允许，请勿转载！
[返回Spark教程首页]

CountVectorizer旨在通过计数来将一个文档转换为向量。当不存在先验字典时，Countvectorizer作为Estimator提取词汇进行训练，并生成一个CountVectorizerModel用于存储相应的词汇向量空间。该模型产生文档关于词语的稀疏表示，其表示可以传递给其他算法，例如LDA。

在CountVectorizerModel的训练过程中，CountVectorizer将根据语料库中的词频排序从高到低进行选择，词汇表的最大含量由vocabsize超参数来指定，超参数minDF，则指定词汇表中的词语至少要在多少个不同文档中出现。

由于Spark2.0起，SQLContext、HiveContext已经不再推荐使用，改以SparkSession代之，故本文中不再使用SQLContext来进行相关的操作，关于SparkSession的具体详情，这里不再赘述，可以参看Spark2.0的官方文档。Spark2.0以上版本的spark-shell在启动时会自动创建一个名为spark的SparkSession对象，当需要手工创建时，SparkSession可以由其伴生对象的builder()方法创建出来，如下代码段所示：

import org.apache.spark.sql.SparkSession
val spark = SparkSession.builder().
            master("local").
            appName("my App Name").
            getOrCreate()

和SQLContext一样，也可以开启RDD的隐式转换：

import spark.implicits._

下文中，我们默认名为spark的SparkSession已经创建。

我们接下来通过一个例子来进行介绍。首先，导入CountVectorizer所需要的包：

import org.apache.spark.ml.feature.{CountVectorizer, CountVectorizerModel}

假设我们有如下的DataFrame，其包含id和words两列，可以看成是一个包含两个文档的迷你语料库。

scala> val df = spark.createDataFrame(Seq(
   |       (0, Array("a", "b", "c")),
   |       (1, Array("a", "b", "b", "c", "a"))
   |     )).toDF("id", "words")
df: org.apache.spark.sql.DataFrame = [id: int, words: array<string>]

随后，通过CountVectorizer设定超参数，训练一个CountVectorizerModel，这里设定词汇表的最大量为3，设定词汇表中的词至少要在2个文档中出现过，以过滤那些偶然出现的词汇。

scala> val cvModel: CountVectorizerModel = new CountVectorizer().
   |       setInputCol("words").
   |       setOutputCol("features").
   |       setVocabSize(3).
   |       setMinDF(2).
   |       fit(df)
cvModel: org.apache.spark.ml.feature.CountVectorizerModel = cntVec_237a080886a2

在训练结束后，可以通过CountVectorizerModel的vocabulary成员获得到模型的词汇表：

scala> cvModel.vocabulary
res7: Array[String] = Array(b, a, c)

从打印结果我们可以看到，词汇表中有“a”，“b”，“c”三个词，且这三个词都在2个文档中出现过（前文设定了minDF为2）。

使用这一模型对DataFrame进行变换，可以得到文档的向量化表示：

scala> cvModel.transform(df).show(false)
+---+---------------+-------------------------+
|id |words          |features                 |
+---+---------------+-------------------------+
|0  |[a, b, c]      |(3,[0,1,2],[1.0,1.0,1.0])|
|1  |[a, b, b, c, a]|(3,[0,1,2],[2.0,2.0,1.0])|
+---+---------------+-------------------------+

和其他Transformer不同，CountVectorizerModel可以通过指定一个先验词汇表来直接生成，如以下例子，直接指定词汇表的成员是“a”，“b”，“c”三个词：

scala> val cvm = new CountVectorizerModel(Array("a", "b", "c")).
   |       setInputCol("words").
   |       setOutputCol("features")
cvm: org.apache.spark.ml.feature.CountVectorizerModel = cntVecModel_c6a17c2befee

scala> cvm.transform(df).select("features").foreach { println }    
[(3,[0,1,2],[1.0,1.0,1.0])]
[(3,[0,1,2],[2.0,2.0,1.0])]

子雨大数据之Spark入门
扫一扫访问本博客

厦大数据库实验室博客