ComputeUtils { public static <K, V, C extends Serializable> 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) { affinityCallWithRetries(ignite, Arrays.asList(datasetCacheName, upstreamCacheName), part -> { Ignite locIgnite = Ignition.localIgnite(); LearningEnvironment env = envBuilder.buildForWorker(part); IgniteCache<K, V> locUpstreamCache = locIgnite.cache(upstreamCacheName); if (isKeepBinary) locUpstreamCache = locUpstreamCache.withKeepBinary(); ScanQuery<K, V> qry = new ScanQuery<>(); qry.setLocal(true); qry.setPartition(part); qry.setFilter(filter); C ctx; UpstreamTransformer transformer = transformerBuilder.build(env); UpstreamTransformer transformerCp = Utils.copy(transformer); long cnt = computeCount(locUpstreamCache, qry, transformer); try (QueryCursor<UpstreamEntry<K, V>> cursor = locUpstreamCache.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 + ']'); ctx = ctxBuilder.build(env, iter, cnt); } IgniteCache<Integer, C> datasetCache = locIgnite.cache(datasetCacheName); datasetCache.put(part, ctx); return part; }, retries, interval, deployingContext); }  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 testInitContext() { 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); ComputeUtils.<Integer, Integer, Integer>initContext( ignite, upstreamCacheName, UpstreamTransformerBuilder.identity(), (k, v) -> true, datasetCacheName, (env, upstream, upstreamSize) -> { assertEquals(1, upstreamSize); UpstreamEntry<Integer, Integer> e = upstream.next(); return e.getKey() + e.getValue(); }, TestUtils.testEnvBuilder(), 0, 0, false, DeployingContext.unitialized() ); assertEquals(1, datasetCache.size()); assertEquals(84, datasetCache.get(0).intValue()); }