spark-mllib-TFIDF实现

TF就是词在一篇文章中的词频，IDF就是逆词频，IFIDF就是两者乘积，常用来表示词在文章中重要性，公式表示为：

官网上给出使用IF-IDF的例子代码：

object TfIdfTest {

  def main(args:Array[String]){

    val conf = new SparkConf().setAppName("TfIdfTest")

    val sc = new SparkContext(conf)

    // Load documents (one per line).

    val documents: RDD[Seq[String]] = sc.textFile("...").map(_.split(" ").toSeq)

    val hashingTF = new HashingTF()

    val tf: RDD[Vector] = hashingTF.transform(documents)

    tf.cache()

    val idf = new IDF().fit(tf)

    val tfidf: RDD[Vector] = idf.transform(tf)

  }

}

可以看到重要的就是三个类， HashingTF，IDF，IDFModel，其中val idf 的类型就是 IDFModel。

首先明确的是，它要求源数据为一篇文章一行

先看 val tf: RDD[Vector] = hashingTF.transform(documents)，将调用 HashingTF类的如下方法

/**

   * Transforms the input document to term frequency vectors.

   */

  def transform[D <: Iterable[_]](dataset: RDD[D]): RDD[Vector] = {

    dataset.map(this.transform)

  }

该方法对每一篇文章调用：

/**

   * Transforms the input document into a sparse term frequency vector.

   */

  def transform(document: Iterable[_]): Vector = {

    val termFrequencies = mutable.HashMap.empty[Int, Double]

    document.foreach { term =>

      val i = indexOf(term)

      termFrequencies.put(i, termFrequencies.getOrElse(i, 0.0) + 1.0)

    }

    Vectors.sparse(numFeatures, termFrequencies.toSeq)

  }

该方法会将词特征映射到一个很大维度的向量中去，每篇文章一个向量，向量长度为 numFeatures，这是 HashingTF类的成员变量，默认为2的20次方。映射时先使用indexOf(term)获取词项term对应的向量位置，该方法如下：

/**

   * Returns the index of the input term.

   */

  def indexOf(term: Any): Int = Utils.nonNegativeMod(term.##, numFeatures)

term.##是得到一个int型的hash值，然后对numFeatures取模，得到位置后，在对应位置上进行加1操作，即：termFrequencies.put(i, termFrequencies.getOrElse(i, 0.0) + 1.0)，最后把转化为一个稀疏的向量。

这样就得到了所有文章的词频向量，一篇文章一个向量。

val idf = new IDF().fit(tf)将调用IDF类的fit方法并返回一个 IDFModel 对象。源码如下：

/**

   * Computes the inverse document frequency.

   * @param dataset an RDD of term frequency vectors

   */

  def fit(dataset: RDD[Vector]): IDFModel = {

    val idf = dataset.treeAggregate(new IDF.DocumentFrequencyAggregator)(

      seqOp = (df, v) => df.add(v),

      combOp = (df1, df2) => df1.merge(df2)

    ).idf()

    new IDFModel(idf)

  }

该方法将完成IDF的所有计算工作。dataset.treeAggregate(zeroValue: U)(seqOp, combOp)的作用是使用方法 seqOp加入新的需计算聚合对象和 combOp聚合合并各自的已聚合结果（因为 treeAggregate是使用mapPartition进行计算的，最后需要合并结果 ） 。该调用返回一个类型为U的对象，在这里就是 IDF.DocumentFrequencyAggregator的对象了，该返回对象再调用方法 idf()完成 IDF 的除法计算，最后通过new IDFModel(idf) 得到一个IDFModel对象。来看看  IDF.DocumentFrequencyAggregator 的源码：

/** Document frequency aggregator. */

  class DocumentFrequencyAggregator extends Serializable {

    /** number of documents */

    private var m = 0L

    /** document frequency vector */

    private var df: BDV[Long] = _

    /** Adds a new document. */

    def add(doc: Vector): this.type = {

      if (isEmpty) {

        df = BDV.zeros(doc.size)

      }

      doc match {

        case sv: SparseVector =>

          val nnz = sv.indices.size

          var k = 0

          while (k < nnz) {

            if (sv.values(k) > 0) {

              df(sv.indices(k)) += 1L

            }

            k += 1

          }

        case dv: DenseVector =>

          val n = dv.size

          var j = 0

          while (j < n) {

            if (dv.values(j) > 0.0) {

              df(j) += 1L

            }

            j += 1

          }

        case other =>

          throw new UnsupportedOperationException(

            s"Only sparse and dense vectors are supported but got ${other.getClass}.")

      }

      m += 1L

      this

    }

    /** Merges another. */

    def merge(other: DocumentFrequencyAggregator): this.type = {

      if (!other.isEmpty) {

        m += other.m

        if (df == null) {

          df = other.df.copy

        } else {

          df += other.df

        }

      }

      this

    }

    private def isEmpty: Boolean = m == 0L

    /** Returns the current IDF vector. */

    def idf(): Vector = {

      if (isEmpty) {

        throw new IllegalStateException("Haven't seen any document yet.")

      }

      val n = df.length

      val inv = new Array[Double](n)

      var j = 0

      while (j < n) {

        inv(j) = math.log((m + 1.0)/ (df(j) + 1.0))

        j += 1

      }

      Vectors.dense(inv)

    }

  }

一个DocumentFrequencyAggregator对象代表了一个partition里面的聚合结果，成员 m 表示该partition里有多少篇文章，df: BDV[Long] 表示词出现的文档频率向量，向量的每个元素表示这个词在多少篇文章中出现过。

看add方法，对每一个加入进来的文档Vector，会循环取出Vector中的term，如果该term对应的位置的值是大于0的，即表示该词在这篇文档中出现了，那么 df(sv.indices(k)) += 1L。

看merge方法，这是两个partition进行聚合的时候使用，即把对应的成员值相加即可。

看idf方法，也只是简单的进行最终的log计算，最后返回结果向量。

得到IDFModel对象后，进行最后的val tfidf: RDD[Vector] = idf.transform(tf)操作，得到最终结果：

/**

   * Transforms term frequency (TF) vectors to TF-IDF vectors.

   * @param dataset an RDD of term frequency vectors

   * @return an RDD of TF-IDF vectors

   */

  def transform(dataset: RDD[Vector]): RDD[Vector] = {

    val bcIdf = dataset.context.broadcast(idf)

    dataset.mapPartitions { iter =>

      val thisIdf = bcIdf.value

      iter.map { v =>

        val n = v.size

        v match {

          case sv: SparseVector =>

            val nnz = sv.indices.size

            val newValues = new Array[Double](nnz)

            var k = 0

            while (k < nnz) {

              newValues(k) = sv.values(k) * thisIdf(sv.indices(k))

              k += 1

            }

            Vectors.sparse(n, sv.indices, newValues)

          case dv: DenseVector =>

            val newValues = new Array[Double](n)

            var j = 0

            while (j < n) {

              newValues(j) = dv.values(j) * thisIdf(j)

              j += 1

            }

            Vectors.dense(newValues)

          case other =>

            throw new UnsupportedOperationException(

              s"Only sparse and dense vectors are supported but got ${other.getClass}.")

        }

      }

    }

  }

先将fit方法中得到的逆文档词频向量idf进行广播，然后进行简单的相乘操作而已。

顺便提一下mllib中的Vector，分为DenseVector和SparseVector， DenseVector就是扎扎实实的顺序存储向量中的每一个元素，SparseVector略有不同，成员size是真正的向量长度，indices是记录所有值不为0的元素的索引，values是记录所有不为0的元素的值，比如一个向量为（0,1,0,0,2,0,0），那么indices则为数组（1,4），values则为数组（1,2），这两个数组必须长度相同。