WeightSampleBatchOp extends BatchOperator<WeightSampleBatchOp> implements WeightSampleParams<WeightSampleBatchOp> { @Override public WeightSampleBatchOp linkFrom(BatchOperator<?>... inputs) { BatchOperator in = checkAndGetFirst(inputs); DataSet<Row> data = in.getDataSet(); int weightIdx = TableUtil.findColIndexWithAssertAndHint(in.getColNames(), getWeightCol()); double ratio = getRatio(); if(getWithReplacement()) { DataSet<Tuple3<Integer, Integer, Double>> weight = data.mapPartition( new RichMapPartitionFunction<Row, Tuple3<Integer, Integer, Double>>() { @Override public void mapPartition(Iterable<Row> values, Collector<Tuple3<Integer, Integer, Double>> out) throws Exception { int cnt = 0; double sum = 0.; int taskId = getRuntimeContext().getIndexOfThisSubtask(); for (Row row : values) { double weight = ((Number)row.getField(weightIdx)).doubleValue(); Preconditions.checkArgument(weight > 0 && !Double.isNaN(weight) && Double.isFinite(weight), "Weight must be positive!"); cnt++; sum += weight; } out.collect(Tuple3.of(taskId, cnt, sum)); } }); DataSet<Tuple2<Integer, double[]>> bounds = weight.reduceGroup( new GroupReduceFunction<Tuple3<Integer, Integer, Double>, Tuple2<Integer, double[]>>() { @Override public void reduce(Iterable<Tuple3<Integer, Integer, Double>> values, Collector<Tuple2<Integer, double[]>> out) throws Exception { List<Tuple3<Integer, Integer, Double>> list = new ArrayList<>(); values.forEach(list::add); list.sort(Comparator.comparingDouble(t -> t.f0)); double[] bounds = new double[list.size() + 1]; int cnt = 0; for(int i = 0; i < list.size(); i++){ bounds[i + 1] = bounds[i] + list.get(i).f2; cnt += list.get(i).f1; } out.collect(Tuple2.of(cnt, bounds)); } }); final Random random = new Random(0); DataSet<Row> res = data .mapPartition(new RandomSelect(random, ratio, weightIdx)) .withBroadcastSet(bounds, BOUNDS); this.setOutput(res, in.getSchema()); }else { DataSet<Tuple2<Double, Row>> weight = data.mapPartition( new RichMapPartitionFunction<Row, Tuple2<Double, Row>>() { @Override public void mapPartition(Iterable<Row> values, Collector<Tuple2<Double, Row>> out) throws Exception { Random random = new Random(getRuntimeContext().getIndexOfThisSubtask()); for (Row row : values) { double weight = ((Number)row.getField(weightIdx)).doubleValue(); Preconditions.checkArgument(weight > 0 && !Double.isNaN(weight) && Double.isFinite(weight), "Weight must be positive!"); double rp = random.nextDouble(); while (rp <= 1e-30) { rp = random.nextDouble(); } out.collect(Tuple2.of(Math.log(rp) / weight, row)); } } }).partitionByRange(0); DataSet<Tuple3<Integer, Integer, Double>> taskCnts = weight.mapPartition( new RichMapPartitionFunction<Tuple2<Double, Row>, Tuple3<Integer, Integer, Double>>() { @Override public void mapPartition(Iterable<Tuple2<Double, Row>> values, Collector<Tuple3<Integer, Integer, Double>> out) throws Exception { int taskId = getRuntimeContext().getIndexOfThisSubtask(); int cnt = 0; double min = Double.MAX_VALUE; for (Tuple2<Double, Row> t : values) { min = Math.min(t.f0, min); cnt++; } out.collect(Tuple3.of(taskId, cnt, min)); } }); DataSet<Row> res = weight.mapPartition(new TopNSelect(ratio)) .withBroadcastSet(taskCnts, COUNT); this.setOutput(res, in.getSchema()); } return this; } WeightSampleBatchOp();  WeightSampleBatchOp(Params params); @Override WeightSampleBatchOp linkFrom(BatchOperator<?>... inputs);  }
@Test public void testWithoutReplacement() throws Exception { Row[] testArray = new Row[] { Row.of("a", 1.3, 1.1), Row.of("b", 2.5, 0.9), Row.of("c", 100., -0.01), Row.of("d", 100., 100.9), Row.of("e", 1.4, 1.1), Row.of("f", 2.2, 0.9), Row.of("g", 100.1, -0.01), Row.of("j", 100., 100.9) }; String[] colnames = new String[] {"id", "weight", "col0"}; MemSourceBatchOp inOp = new MemSourceBatchOp(Arrays.asList(testArray), colnames); WeightSampleBatchOp sampleBatchOp = new WeightSampleBatchOp() .setWeightCol("weight") .setRatio(0.3) .linkFrom(inOp); Assert.assertEquals(sampleBatchOp.count(), 2); }
@Test public void testWithReplacement() throws Exception { Row[] testArray = new Row[] { Row.of("a", 1.3, 1.1), Row.of("b", 2.5, 0.9), Row.of("c", 100., -0.01), Row.of("d", 100., 100.9), Row.of("e", 1.4, 1.1), Row.of("f", 2.2, 0.9), Row.of("g", 100.1, -0.01), Row.of("j", 100., 100.9) }; String[] colnames = new String[] {"id", "weight", "col0"}; MemSourceBatchOp inOp = new MemSourceBatchOp(Arrays.asList(testArray), colnames); WeightSampleBatchOp sampleBatchOp = new WeightSampleBatchOp() .setWeightCol("weight") .setRatio(0.4) .setWithReplacement(true) .linkFrom(inOp); Assert.assertEquals(sampleBatchOp.count(), 3); }