SplitBatchOp extends BatchOperator<SplitBatchOp> implements SplitParams<SplitBatchOp> { @Override public SplitBatchOp linkFrom(BatchOperator<?>... inputs) { BatchOperator<?> in = checkAndGetFirst(inputs); final double fraction = getFraction(); if (fraction < 0. || fraction > 1.0) { throw new RuntimeException("invalid fraction " + fraction); } DataSet<Row> rows = in.getDataSet(); DataSet<Tuple2<Integer, Long>> countsPerPartition = DataSetUtils.countElementsPerPartition(rows); DataSet<long[]> numPickedPerPartition = countsPerPartition .mapPartition(new CountInPartition(fraction)) .setParallelism(1) .name("decide_count_of_each_partition"); DataSet<Row> out = rows .mapPartition(new PickInPartition()) .withBroadcastSet(numPickedPerPartition, "counts") .name("pick_in_each_partition"); this.setOutput(out, in.getSchema()); this.setSideOutputTables(new Table[]{in.getOutputTable().minusAll(this.getOutputTable())}); return this; } SplitBatchOp();  SplitBatchOp(Params params);  SplitBatchOp(double fraction); @Override SplitBatchOp linkFrom(BatchOperator<?>... inputs);  }
@Test public void split() throws Exception { BatchOperator data = Iris.getBatchData(); data = new AppendIdBatchOp().linkFrom(data); BatchOperator spliter = new SplitBatchOp().setFraction(0.1); BatchOperator left = spliter.linkFrom(data); BatchOperator right = spliter.getSideOutput(0); Assert.assertEquals(left.count(), 15); Assert.assertEquals(right.count(), 150 - 15); Assert.assertEquals(new UnionBatchOp().linkFrom(left, right).count(), 150); }