使用Spark和Java进行分层抽样

时间:2017-01-16 09:10:53

标签: java apache-spark machine-learning apache-spark-mllib

我想确保我正在对我的数据进行分层抽样培训。

按照here解释,Spark 2.1和早期版本似乎支持JavaPairRDD.sampleByKey(...)JavaPairRDD.sampleByKeyExact(...)

但是:我的数据存储在Dataset<Row>,而不是JavaPairRDD。第一列是标签,所有其他都是功能(从libsvm格式的文件导入)。

获取数据集实例的分层样本的最简单方法是什么?最后再次有Dataset<Row>

在某种程度上,此问题与Dealing with unbalanced datasets in Spark MLlib 相关。

这个possible duplicate根本没有提到Dataset<Row>,也没有提到Java。它没有回答我的问题。

1 个答案:

答案 0 :(得分:4)

好的,既然the question here的答案实际上并不适用于 Java ,我已经用 Java 重写了它。

推理仍然是一样的想法。我们仍在使用sampleByKeyExact。目前没有开箱即用的奇迹功能( spark 2.1.0

所以你走了:

package org.awesomespark.examples;

import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.PairFunction;
import org.apache.spark.sql.*;
import scala.Tuple2;

import java.util.Map;

public class StratifiedDatasets {
    public static void main(String[] args) {
        SparkSession spark = SparkSession.builder()
                .appName("Stratified Datasets")
                .getOrCreate();

        Dataset<Row> data = spark.read().format("libsvm").load("sample_libsvm_data.txt");

        JavaPairRDD<Double, Row> rdd = data.toJavaRDD().keyBy(x -> x.getDouble(0));
        Map<Double, Double> fractions = rdd.map(Tuple2::_1)
                .distinct()
                .mapToPair((PairFunction<Double, Double, Double>) (Double x) -> new Tuple2(x, 0.8))
                .collectAsMap();

        JavaRDD<Row> sampledRDD = rdd.sampleByKeyExact(false, fractions, 2L).values();
        Dataset<Row> sampledData = spark.createDataFrame(sampledRDD, data.schema());

        sampledData.show();
        sampledData.printSchema();
    }
}

现在让我们打包并提交:

$ sbt package
[...]
// [success] Total time: 2 s, completed Jan 16, 2017 1:45:51 PM

$ spark-submit --class org.awesomespark.examples.StratifiedDatasets target/scala-2.10/java-stratified-dataset_2.10-1.0.jar 
[...]

// +-----+--------------------+
// |label|            features|
// +-----+--------------------+
// |  0.0|(692,[127,128,129...|
// |  1.0|(692,[158,159,160...|
// |  1.0|(692,[124,125,126...|
// |  1.0|(692,[152,153,154...|
// |  1.0|(692,[151,152,153...|
// |  0.0|(692,[129,130,131...|
// |  1.0|(692,[99,100,101,...|
// |  0.0|(692,[154,155,156...|
// |  0.0|(692,[127,128,129...|
// |  1.0|(692,[154,155,156...|
// |  0.0|(692,[151,152,153...|
// |  1.0|(692,[129,130,131...|
// |  0.0|(692,[154,155,156...|
// |  1.0|(692,[150,151,152...|
// |  0.0|(692,[124,125,126...|
// |  0.0|(692,[152,153,154...|
// |  1.0|(692,[97,98,99,12...|
// |  1.0|(692,[124,125,126...|
// |  1.0|(692,[156,157,158...|
// |  1.0|(692,[127,128,129...|
// +-----+--------------------+
// only showing top 20 rows

// root
//  |-- label: double (nullable = true)
//  |-- features: vector (nullable = true)

对于 python 用户,您还可以查看我的回答Stratified sampling with pyspark