class 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()) {
            //TaskId, count, weightSum
            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);

            //taskId, count, minValue
            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;
    }

    public  WeightSampleBatchOp();
    public  WeightSampleBatchOp(Params params);

}