spark mllib IDF源码解析

通常在文本向量化的过程中，通常用的比较多的就是tf-idf、word2vec、CountVectorizer的这几个方法，前面的博客有分析过HashingTF和CountVectorizer方法，今天来分析IDF，
通常IDF和HashingTF一期联合使用。
1、IDF计算公式：

计算公式：idf = log((m + 1) / (d(t) + 1))
m : 代表语料库文档数量 t ：包含该词条的文档数量
m+1 和 dt(t) + 1 如果该词语不在语料库中，就会导致被除数为零，因此一般情况下使用 dt(t) + 1

2、参数

val minDocFreq 默认为0即最小的文档数

3、装箱

def fit(dataset: RDD[Vector]): IDFModel = {
    // 调用treeAggregate方法 完成数据聚合
    val idf = dataset.treeAggregate(new IDF.DocumentFrequencyAggregator(
          minDocFreq = minDocFreq))(
      // 在每个partition里面初始化一个DocumentFrequencyAggregator
      // 将每个文档调用其add方法 合并到一df对象
      seqOp = (df, v) => df.add(v),
      // 分区间DocumentFrequencyAggregator对象调用自身merge方法实现合并
      // 最终在driver端完成最终的合并
      combOp = (df1, df2) => df1.merge(df2)
    ).idf()
    new IDFModel(idf)
  }

4、查看DocumentFrequencyAggregator 的add方法和merge方法

 def add(doc: Vector): this.type = {
      if (isEmpty) {
        df = BDV.zeros(doc.size)
      }

      // 模式匹配稀疏向量和稠密向量
      // 统计词元出现的文档数
      doc match {
        case SparseVector(size, indices, values) =>
          val nnz = indices.length
          var k = 0
          while (k < nnz) {
            if (values(k) > 0) {
              df(indices(k)) += 1L
            }
            k += 1
          }
        case DenseVector(values) =>
          val n = values.length
          var j = 0
          while (j < n) {
            if (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
    }
    // 完成合并
def merge(other: DocumentFrequencyAggregator): this.type = {
      if (!other.isEmpty) {
        m += other.m
        if (df == null) {
          df = other.df.copy
        } else {
          df += other.df
        }
      }
      this
    }

5、计算IDF

 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) {
         // 如果词元没有出现，则idf默认为0，这样会造成tf-idf也为0
         // 为了平滑数据，引入+1
        if (df(j) >= minDocFreq) {
          inv(j) = math.log((m + 1.0) / (df(j) + 1.0))
        }
        j += 1
      }
      Vectors.dense(inv)
    }

到此IDF计算完成，其实就算过程很简单；
Tips：
当你的模型在训练的时候会生成IDFModel方法，在模型训练完成后需要对IDFModel保存，
在后续使用模型预估数据的过程，需要用IDFModel去转向量，这样能最大限度保证输入数据的一致性。

后续会详细分析treeAggregate和Aggregate方法。