LocalDatasetBuilder implements DatasetBuilder<K, V> { @Override public <C extends Serializable, D extends AutoCloseable> LocalDataset<C, D> build( LearningEnvironmentBuilder envBuilder, PartitionContextBuilder<K, V, C> partCtxBuilder, PartitionDataBuilder<K, V, C, D> partDataBuilder, LearningEnvironment learningEnvironment) { List<C> ctxList = new ArrayList<>(); List<D> dataList = new ArrayList<>(); List<UpstreamEntry<K, V>> entriesList = new ArrayList<>(); upstreamMap .entrySet() .stream() .filter(en -> filter.apply(en.getKey(), en.getValue())) .map(en -> new UpstreamEntry<>(en.getKey(), en.getValue())) .forEach(entriesList::add); int partSize = Math.max(1, entriesList.size() / partitions); Iterator<UpstreamEntry<K, V>> firstKeysIter = entriesList.iterator(); Iterator<UpstreamEntry<K, V>> secondKeysIter = entriesList.iterator(); Iterator<UpstreamEntry<K, V>> thirdKeysIter = entriesList.iterator(); int ptr = 0; List<LearningEnvironment> envs = IntStream.range(0, partitions).boxed().map(envBuilder::buildForWorker) .collect(Collectors.toList()); for (int part = 0; part < partitions; part++) { int cntBeforeTransform = part == partitions - 1 ? entriesList.size() - ptr : Math.min(partSize, entriesList.size() - ptr); LearningEnvironment env = envs.get(part); UpstreamTransformer transformer1 = upstreamTransformerBuilder.build(env); UpstreamTransformer transformer2 = Utils.copy(transformer1); UpstreamTransformer transformer3 = Utils.copy(transformer1); int cnt = (int)transformer1.transform(Utils.asStream(new IteratorWindow<>(thirdKeysIter, k -> k, cntBeforeTransform))).count(); Iterator<UpstreamEntry> iter = transformer2.transform(Utils.asStream(new IteratorWindow<>(firstKeysIter, k -> k, cntBeforeTransform)).map(x -> (UpstreamEntry)x)).iterator(); Iterator<UpstreamEntry<K, V>> convertedBack = Utils.asStream(iter).map(x -> (UpstreamEntry<K, V>)x).iterator(); C ctx = cntBeforeTransform > 0 ? partCtxBuilder.build(env, convertedBack, cnt) : null; Iterator<UpstreamEntry> iter1 = transformer3.transform( Utils.asStream(new IteratorWindow<>(secondKeysIter, k -> k, cntBeforeTransform))).iterator(); Iterator<UpstreamEntry<K, V>> convertedBack1 = Utils.asStream(iter1).map(x -> (UpstreamEntry<K, V>)x).iterator(); D data = cntBeforeTransform > 0 ? partDataBuilder.build( env, convertedBack1, cnt, ctx ) : null; ctxList.add(ctx); dataList.add(data); ptr += cntBeforeTransform; } return new LocalDataset<>(envs, ctxList, dataList); } LocalDatasetBuilder(Map<K, V> upstreamMap, int partitions);  LocalDatasetBuilder(Map<K, V> upstreamMap,
        IgniteBiPredicate<K, V> filter,
        int partitions,
        UpstreamTransformerBuilder upstreamTransformerBuilder);  LocalDatasetBuilder(Map<K, V> upstreamMap,
        IgniteBiPredicate<K, V> filter,
        int partitions); @Override LocalDataset<C, D> build(
        LearningEnvironmentBuilder envBuilder,
        PartitionContextBuilder<K, V, C> partCtxBuilder, PartitionDataBuilder<K, V, C, D> partDataBuilder,
        LearningEnvironment learningEnvironment); @Override DatasetBuilder<K, V> withUpstreamTransformer(UpstreamTransformerBuilder builder); @Override DatasetBuilder<K, V> withFilter(IgniteBiPredicate<K, V> filterToAdd);  }
@Test public void testBuild() { Map<Integer, Integer> data = new HashMap<>(); for (int i = 0; i < 100; i++) data.put(i, i); LocalDatasetBuilder<Integer, Integer> builder = new LocalDatasetBuilder<>(data, 10); LocalDataset<Serializable, TestPartitionData> dataset = buildDataset(builder); assertEquals(10, dataset.getCtx().size()); assertEquals(10, dataset.getData().size()); AtomicLong cnt = new AtomicLong(); dataset.compute((partData, env) -> { cnt.incrementAndGet(); int[] arr = partData.data; assertEquals(10, arr.length); for (int i = 0; i < 10; i++) assertEquals(env.partition() * 10 + i, arr[i]); }); assertEquals(10, cnt.intValue()); }