将元数据附加到Spark中的矢量列

Question

上下文: 我有一个包含两列的数据框:标签和功能.

org.apache.spark.sql.DataFrame = [label: int, features: vector]

其中features是使用VectorAssembler构建的数值类型的mllib.linalg.VectorUDT.

问题: 有没有办法为特征向量分配模式？我想跟踪每个功能的名称.

到目前为止尝试过:

val defaultAttr = NumericAttribute.defaultAttr
val attrs = Array("feat1", "feat2", "feat3").map(defaultAttr.withName)
val attrGroup = new AttributeGroup("userFeatures", attrs.asInstanceOf[Array[Attribute]])

---

scala> attrGroup.toMetadata 
res197: org.apache.spark.sql.types.Metadata = {"ml_attr":{"attrs":{"numeric":[{"idx":0,"name":"f1"},{"idx":1,"name":"f2"},{"idx":2,"name":"f3"}]},"num_attrs":3}}

但不确定如何将其应用于现有数据框.

Answer 1

至少有两种选择:

1. 在现有的DataFrame你可以使用as带metadata参数的方法:

   import org.apache.spark.ml.attribute._
   
   val rdd = sc.parallelize(Seq(
     (1, Vectors.dense(1.0, 2.0, 3.0))
   ))
   val df = rdd.toDF("label", "features")
   
   df.withColumn("features", $"features".as("_", attrGroup.toMetadata))
   

2. 创建新DataFrame转换AttributeGroup toStructField并将其用作给定列的架构时:

   import org.apache.spark.sql.types.{StructType, StructField, IntegerType}
   
   val schema = StructType(Array(
     StructField("label", IntegerType, false),
     attrGroup.toStructField()
   ))
   
   spark.createDataFrame(
     rdd.map(row => Row.fromSeq(row.productIterator.toSeq)),
     schema)
   

如果使用VectorAssembler列元数据创建了向量列,则应该已经附加了描述父列的列.

import org.apache.spark.ml.feature.VectorAssembler

val raw = sc.parallelize(Seq(
  (1, 1.0, 2.0, 3.0)
)).toDF("id", "feat1", "feat2", "feat3")

val assembler = new VectorAssembler()
  .setInputCols(Array("feat1", "feat2", "feat3"))
  .setOutputCol("features")

val dfWithMeta = assembler.transform(raw).select($"id", $"features")
dfWithMeta.schema.fields(1).metadata

// org.apache.spark.sql.types.Metadata = {"ml_attr":{"attrs":{"numeric":[
//   {"idx":0,"name":"feat1"},{"idx":1,"name":"feat2"},
//   {"idx":2,"name":"feat3"}]},"num_attrs":3}

使用点语法(如$features.feat1)不能直接访问矢量字段,但可以使用专门的工具,例如VectorSlicer:

import org.apache.spark.ml.feature.VectorSlicer

val slicer = new VectorSlicer()
  .setInputCol("features")
  .setOutputCol("featuresSubset")
  .setNames(Array("feat1", "feat3"))

slicer.transform(dfWithMeta).show
// +---+-------------+--------------+
// | id|     features|featuresSubset|
// +---+-------------+--------------+
// |  1|[1.0,2.0,3.0]|     [1.0,3.0]|
// +---+-------------+--------------+

对于PySpark,请参阅如何将列声明为DataFrame 中的分类功能,以便在ml中使用