预排序输入的Spark特征向量变换

HDFS上的制表符分隔文件中有一些数据,如下所示:

label | user_id | feature
------------------------------
  pos | 111     | www.abc.com
  pos | 111     | www.xyz.com
  pos | 111     | Firefox
  pos | 222     | www.example.com
  pos | 222     | www.xyz.com
  pos | 222     | IE
  neg | 333     | www.jkl.com
  neg | 333     | www.xyz.com
  neg | 333     | Chrome

我需要对其进行转换,为每个user_id创建一个特征向量,以训练 org.apache.spark.ml.classification.NaiveBayes 模型

我目前的做法基本如下:

1. 将原始数据加载到DataFrame
2. 使用StringIndexer索引功能
3. 转到RDD和Group by user_id,并将特征索引映射到稀疏向量中。

关键是 数据已按user_id预先排序 。利用这一点的最佳方法是什么?想想可能会发生多少不必要的工作,我很痛苦。

如果一点代码有助于理解我当前的方法,下面是地图的要点:

val featurization = (vals: (String,Iterable[Row])) => {
  // create a Seq of all the feature indices
  // Note: the indexing was done in a previous step not shown
  val seq = vals._2.map(x => (x.getDouble(1).toInt,1.0D)).toSeq

  // create the sparse vector
  val featureVector = Vectors.sparse(maxIndex, seq)

  // convert the string label into a Double
  val label = if (vals._2.head.getString(2) == "pos") 1.0 else 0.0

  (label, vals._1, featureVector)
}

d.rdd
  .groupBy(_.getString(1))
  .map(featurization)
  .toDF("label","user_id","features")