ComputeUtils { public static <K, V, C extends Serializable, D extends AutoCloseable> D getData( Ignite ignite, String upstreamCacheName, IgniteBiPredicate<K, V> filter, UpstreamTransformerBuilder transformerBuilder, String datasetCacheName, UUID datasetId, PartitionDataBuilder<K, V, C, D> partDataBuilder, LearningEnvironment env, boolean isKeepBinary) { PartitionDataStorage dataStorage = (PartitionDataStorage)ignite .cluster() .nodeLocalMap() .computeIfAbsent(String.format(DATA_STORAGE_KEY_TEMPLATE, datasetId), key -> new PartitionDataStorage(env.dataTtl())); final int part = env.partition(); return dataStorage.computeDataIfAbsent(part, () -> { IgniteCache<Integer, C> learningCtxCache = ignite.cache(datasetCacheName); C ctx = learningCtxCache.get(part); IgniteCache<K, V> upstreamCache = ignite.cache(upstreamCacheName); if (isKeepBinary) upstreamCache = upstreamCache.withKeepBinary(); ScanQuery<K, V> qry = new ScanQuery<>(); qry.setLocal(true); qry.setPartition(part); qry.setFilter(filter); UpstreamTransformer transformer = transformerBuilder.build(env); UpstreamTransformer transformerCp = Utils.copy(transformer); long cnt = computeCount(upstreamCache, qry, transformer); if (cnt > 0) { try (QueryCursor<UpstreamEntry<K, V>> cursor = upstreamCache.query(qry, e -> new UpstreamEntry<>(e.getKey(), e.getValue()))) { Iterator<UpstreamEntry<K, V>> it = cursor.iterator(); Stream<UpstreamEntry> transformedStream = transformerCp.transform(Utils.asStream(it, cnt).map(x -> (UpstreamEntry)x)); it = Utils.asStream(transformedStream.iterator()).map(x -> (UpstreamEntry<K, V>)x).iterator(); Iterator<UpstreamEntry<K, V>> iter = new IteratorWithConcurrentModificationChecker<>(it, cnt, "Cache expected to be not modified during dataset data building [partition=" + part + ']'); return partDataBuilder.build(env, iter, cnt, ctx); } } return null; }); }  static Collection<R> affinityCallWithRetries(Ignite ignite, Collection<String> cacheNames,
        IgniteFunction<Integer, R> fun, int retries, int interval, DeployingContext deployingContext); static Collection<R> affinityCallWithRetries(Ignite ignite, Collection<String> cacheNames,
        IgniteFunction<Integer, R> fun, int retries, DeployingContext deployingContext); static LearningEnvironment getLearningEnvironment(Ignite ignite,
        UUID datasetId,
        int part,
        LearningEnvironmentBuilder envBuilder); static D getData(
        Ignite ignite,
        String upstreamCacheName, IgniteBiPredicate<K, V> filter,
        UpstreamTransformerBuilder transformerBuilder,
        String datasetCacheName, UUID datasetId,
        PartitionDataBuilder<K, V, C, D> partDataBuilder,
        LearningEnvironment env,
        boolean isKeepBinary); static void removeData(Ignite ignite, UUID datasetId); static void removeLearningEnv(Ignite ignite, UUID datasetId); static void initContext(
        Ignite ignite,
        String upstreamCacheName,
        UpstreamTransformerBuilder transformerBuilder,
        IgniteBiPredicate<K, V> filter,
        String datasetCacheName,
        PartitionContextBuilder<K, V, C> ctxBuilder,
        LearningEnvironmentBuilder envBuilder,
        int retries,
        int interval,
        boolean isKeepBinary,
        DeployingContext deployingContext); static C getContext(Ignite ignite, String datasetCacheName, int part); static void saveContext(Ignite ignite, String datasetCacheName, int part, C ctx);  }
@Test public void testGetData() { ClusterNode node = grid(1).cluster().localNode(); String upstreamCacheName = "CACHE_1_" + UUID.randomUUID(); String datasetCacheName = "CACHE_2_" + UUID.randomUUID(); CacheConfiguration<Integer, Integer> upstreamCacheConfiguration = new CacheConfiguration<>(); upstreamCacheConfiguration.setName(upstreamCacheName); upstreamCacheConfiguration.setAffinity(new TestAffinityFunction(node)); IgniteCache<Integer, Integer> upstreamCache = ignite.createCache(upstreamCacheConfiguration); CacheConfiguration<Integer, Integer> datasetCacheConfiguration = new CacheConfiguration<>(); datasetCacheConfiguration.setName(datasetCacheName); datasetCacheConfiguration.setAffinity(new TestAffinityFunction(node)); IgniteCache<Integer, Integer> datasetCache = ignite.createCache(datasetCacheConfiguration); upstreamCache.put(42, 42); datasetCache.put(0, 0); UUID datasetId = UUID.randomUUID(); IgniteAtomicLong cnt = ignite.atomicLong("CNT_" + datasetId, 0, true); for (int i = 0; i < 10; i++) { Collection<TestPartitionData> data = ComputeUtils.affinityCallWithRetries( ignite, Arrays.asList(datasetCacheName, upstreamCacheName), part -> ComputeUtils.<Integer, Integer, Serializable, TestPartitionData>getData( ignite, upstreamCacheName, (k, v) -> true, UpstreamTransformerBuilder.identity(), datasetCacheName, datasetId, (env, upstream, upstreamSize, ctx) -> { cnt.incrementAndGet(); assertEquals(1, upstreamSize); UpstreamEntry<Integer, Integer> e = upstream.next(); return new TestPartitionData(e.getKey() + e.getValue()); }, TestUtils.testEnvBuilder().buildForWorker(part), false ), 0, DeployingContext.unitialized() ); assertEquals(1, data.size()); TestPartitionData dataElement = data.iterator().next(); assertEquals(84, dataElement.val.intValue()); } assertEquals(1, cnt.get()); }