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方法。