zipWithIndex на Apache Flink

Теперь это часть версии 0.10-SNAPSHOT Apache Flink. Примеры для zipWithIndex(in) и zipWithUniqueId(in) доступны в официальной документации Flink < / а>.

Вот простая реализация функции:

public class ZipWithIndex {

public static void main(String[] args) throws Exception {

    ExecutionEnvironment ee = ExecutionEnvironment.getExecutionEnvironment();

    DataSet<String> in = ee.readTextFile("/home/robert/flink-workdir/debug/input");

    // count elements in each partition
    DataSet<Tuple2<Integer, Long>> counts = in.mapPartition(new RichMapPartitionFunction<String, Tuple2<Integer, Long>>() {
        @Override
        public void mapPartition(Iterable<String> values, Collector<Tuple2<Integer, Long>> out) throws Exception {
            long cnt = 0;
            for (String v : values) {
                cnt++;
            }
            out.collect(new Tuple2<Integer, Long>(getRuntimeContext().getIndexOfThisSubtask(), cnt));
        }
    });

    DataSet<Tuple2<Long, String>> result = in.mapPartition(new RichMapPartitionFunction<String, Tuple2<Long, String>>() {
        long start = 0;

        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            List<Tuple2<Integer, Long>> offsets = getRuntimeContext().getBroadcastVariable("counts");
            Collections.sort(offsets, new Comparator<Tuple2<Integer, Long>>() {
                @Override
                public int compare(Tuple2<Integer, Long> o1, Tuple2<Integer, Long> o2) {
                    return ZipWithIndex.compare(o1.f0, o2.f0);
                }
            });
            for(int i = 0; i < getRuntimeContext().getIndexOfThisSubtask(); i++) {
                start += offsets.get(i).f1;
            }
        }

        @Override
        public void mapPartition(Iterable<String> values, Collector<Tuple2<Long, String>> out) throws Exception {
            for(String v: values) {
                out.collect(new Tuple2<Long, String>(start++, v));
            }
        }
    }).withBroadcastSet(counts, "counts");
    result.print();

}

public static int compare(int x, int y) {
    return (x < y) ? -1 : ((x == y) ? 0 : 1);
}
}

Вот как это работает: я использую первую mapPartition() операцию, чтобы просмотреть все элементы в разделах, чтобы подсчитать, сколько элементов там. Мне нужно знать количество элементов в каждом разделе, чтобы правильно установить смещения при назначении идентификаторов элементам. Результатом первого mapPartition является DataSet, содержащий сопоставления. Я транслирую этот DataSet всем вторым mapPartition() операторам, которые будут назначать идентификаторы элементам из ввода. В open() методе второго mapPartition() я вычисляю смещение для каждого раздела.

Я, вероятно, собираюсь внести код в Flink (после обсуждения с другими коммиттерами).