至少有两种选择:
在现有的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))
当您创建新的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,请参阅如何将 Column 声明为 DataFrame 中的分类特征以在 ml 中使用