使用createDataFrame创建Spark矢量列

时间:2018-05-13 22:11:33

标签: scala apache-spark dataframe vector spark-dataframe

我可以使用带有toDF方法的矢量列创建Spark DataFrame。

val dataset = Seq((1.0, org.apache.spark.ml.linalg.Vectors.dense(0.0, 10.0, 0.5))).toDF("id", "userFeatures")

scala> dataset.printSchema()
root
 |-- id: double (nullable = false)
 |-- userFeatures: vector (nullable = true)


scala> dataset.schema
res5: org.apache.spark.sql.types.StructType = StructType(StructField(id,DoubleType,false), StructField(userFeatures,org.apache.spark.ml.linalg.VectorUDT@3bfc3ba7,true))

我不确定如何使用createDataFrame方法创建矢量列。 org.apache.spark.sql.types中没有VectorType类型。

这不起作用:

val rows = spark.sparkContext.parallelize(
  List(
    Row(1.0, Vectors.dense(1.0, 2.0))
  )
)

val schema = List(
  StructField("id", DoubleType, true),
  StructField("features", new org.apache.spark.ml.linalg.VectorUDT, true)
)

val df = spark.createDataFrame(
  rows,
  StructType(schema)
)

df.show()
df.printSchema()

1 个答案:

答案 0 :(得分:1)

要使用createDataFrame创建Spark矢量列,您可以使用以下代码:

val rows = spark.sparkContext.parallelize(
  List(
    Row(1.0, org.apache.spark.mllib.linalg.Vectors.dense(1.0, 2.0))
  )
)

val schema = List(
  StructField("id", DoubleType, true),
  StructField("features", new org.apache.spark.mllib.linalg.VectorUDT, true)
)

val df = spark.createDataFrame(
  rows,
  StructType(schema)
)

df.show()
+---+---------+
| id| features|
+---+---------+
|1.0|[1.0,2.0]|
+---+---------+

df.printSchema()
root
 |-- id: double (nullable = true)
 |-- features: vector (nullable = true)

实际问题是不兼容的类型org.apache.spark.ml.linalg.Vectors.dense,它不是vector的架构的有效外部类型。因此,我们必须切换到mllib包而不是ml包。

我希望它有所帮助!

注意:我使用的是Spark v2.3.0。此外,无法在包VectorUDT中访问包linalg中的类org.apache.spark.ml.linalg

供参考 - https://github.com/apache/spark/tree/master/mllib/src/main/scala/org/apache/spark/mllib