在java中将csv值转换为sparkDataframe中的向量

voj3qocg  于 2021-06-01  发布在  Hadoop
关注(0)|答案(1)|浏览(443)

我有一个有两列的csv文件

id, features

id列是一个字符串,features列是一个逗号分隔的机器学习算法的特征值列表,即“[1,4,5]”我基本上只需要对值调用vectors.parse()就可以得到一个向量,但我不想先转换为rdd。
我想把它放到spark数据框中,其中features列是 org.apache.spark.mllib.linalg.Vector 我正在用databricks csvapi将其读入一个dataframe,并尝试将features列转换为一个向量。
有人知道如何在java中做到这一点吗?

zlwx9yxi

zlwx9yxi1#

我找到了一种用自定义项的方法。还有其他方法吗?

HashMap<String, String> options = new HashMap<String, String>();
  options.put("header", "true");
  String input= args[0];

  sqlc.udf().register("toVector", new UDF1<String, Vector>() {
     @Override
     public Vector call(String t1) throws Exception {
        return Vectors.parse(t1);
     }
  }, new VectorUDT());

  StructField[] fields = {new StructField("id",DataTypes.StringType,false, Metadata.empty()) , new StructField("features", DataTypes.StringType, false, Metadata.empty())};
  StructType schema = new StructType(fields);

  DataFrame df = sqlc.read().format("com.databricks.spark.csv").schema(schema).options(options).load(input);

  df = df.withColumn("features", functions.callUDF("toVector", df.col("features")));

相关问题