CacheBasedDatasetBuilder implements DatasetBuilder<K, V> { @SuppressWarnings("unchecked") @Override public <C extends Serializable, D extends AutoCloseable> CacheBasedDataset<K, V, C, D> build( LearningEnvironmentBuilder envBuilder, PartitionContextBuilder<K, V, C> partCtxBuilder, PartitionDataBuilder<K, V, C, D> partDataBuilder, LearningEnvironment localLearningEnv) { UUID datasetId = UUID.randomUUID(); CacheConfiguration<K, V> upstreamCacheConfiguration = upstreamCache.getConfiguration(CacheConfiguration.class); AffinityFunction upstreamCacheAffinity = upstreamCacheConfiguration.getAffinity(); CacheConfiguration<Integer, C> datasetCacheConfiguration = new CacheConfiguration<>(); datasetCacheConfiguration.setName(String.format(DATASET_CACHE_TEMPLATE, upstreamCache.getName(), datasetId)); datasetCacheConfiguration.setAffinity(new DatasetAffinityFunctionWrapper(upstreamCacheAffinity)); IgniteCache<Integer, C> datasetCache = ignite.createCache(datasetCacheConfiguration); ComputeUtils.initContext( ignite, upstreamCache.getName(), transformerBuilder, filter, datasetCache.getName(), partCtxBuilder, envBuilder, retries, RETRY_INTERVAL, upstreamKeepBinary, localLearningEnv.deployingContext() ); return new CacheBasedDataset<>(ignite, upstreamCache, filter, transformerBuilder, datasetCache, envBuilder, partDataBuilder, datasetId, upstreamKeepBinary, localLearningEnv, retries); } CacheBasedDatasetBuilder(Ignite ignite, IgniteCache<K, V> upstreamCache);  CacheBasedDatasetBuilder(Ignite ignite, IgniteCache<K, V> upstreamCache, IgniteBiPredicate<K, V> filter);  CacheBasedDatasetBuilder(Ignite ignite,
        IgniteCache<K, V> upstreamCache,
        IgniteBiPredicate<K, V> filter,
        UpstreamTransformerBuilder transformerBuilder);  CacheBasedDatasetBuilder(Ignite ignite,
        IgniteCache<K, V> upstreamCache,
        IgniteBiPredicate<K, V> filter,
        UpstreamTransformerBuilder transformerBuilder,
        Boolean isKeepBinary,
        int retries); @SuppressWarnings("unchecked") @Override CacheBasedDataset<K, V, C, D> build(
        LearningEnvironmentBuilder envBuilder,
        PartitionContextBuilder<K, V, C> partCtxBuilder,
        PartitionDataBuilder<K, V, C, D> partDataBuilder,
        LearningEnvironment localLearningEnv); @Override DatasetBuilder<K, V> withUpstreamTransformer(UpstreamTransformerBuilder builder); @Override DatasetBuilder<K, V> withFilter(IgniteBiPredicate<K, V> filterToAdd); CacheBasedDatasetBuilder<K, V> withKeepBinary(boolean isKeepBinary); CacheBasedDatasetBuilder<K, V> withRetriesNumber(int retries); static final int DEFAULT_NUMBER_OF_RETRIES; }
@Test public void testBuild() { IgniteCache<Integer, String> upstreamCache = createTestCache(100, 10); CacheBasedDatasetBuilder<Integer, String> builder = new CacheBasedDatasetBuilder<>(ignite, upstreamCache); CacheBasedDataset<Integer, String, Long, AutoCloseable> dataset = builder.build( TestUtils.testEnvBuilder(), (env, upstream, upstreamSize) -> upstreamSize, (env, upstream, upstreamSize, ctx) -> null, TestUtils.testEnvBuilder().buildForTrainer() ); Affinity<Integer> upstreamAffinity = ignite.affinity(upstreamCache.getName()); Affinity<Integer> datasetAffinity = ignite.affinity(dataset.getDatasetCache().getName()); int upstreamPartitions = upstreamAffinity.partitions(); int datasetPartitions = datasetAffinity.partitions(); assertEquals(upstreamPartitions, datasetPartitions); for (int part = 0; part < upstreamPartitions; part++) { Collection<ClusterNode> upstreamPartNodes = upstreamAffinity.mapPartitionToPrimaryAndBackups(part); Collection<ClusterNode> datasetPartNodes = datasetAffinity.mapPartitionToPrimaryAndBackups(part); assertEqualsCollections(upstreamPartNodes, datasetPartNodes); } }
@Test public void testBuildWithPredicate() { CacheConfiguration<Integer, Integer> upstreamCacheConfiguration = new CacheConfiguration<>(); upstreamCacheConfiguration.setAffinity(new RendezvousAffinityFunction(false, 1)); upstreamCacheConfiguration.setName(UUID.randomUUID().toString()); IgniteCache<Integer, Integer> upstreamCache = ignite.createCache(upstreamCacheConfiguration); upstreamCache.put(1, 1); upstreamCache.put(2, 2); CacheBasedDatasetBuilder<Integer, Integer> builder = new CacheBasedDatasetBuilder<>( ignite, upstreamCache, (k, v) -> k % 2 == 0 ); CacheBasedDataset<Integer, Integer, Long, AutoCloseable> dataset = builder.build( TestUtils.testEnvBuilder(), (env, upstream, upstreamSize) -> { UpstreamEntry<Integer, Integer> entry = upstream.next(); assertEquals(Integer.valueOf(2), entry.getKey()); assertEquals(Integer.valueOf(2), entry.getValue()); assertFalse(upstream.hasNext()); return 0L; }, (env, upstream, upstreamSize, ctx) -> { UpstreamEntry<Integer, Integer> entry = upstream.next(); assertEquals(Integer.valueOf(2), entry.getKey()); assertEquals(Integer.valueOf(2), entry.getValue()); assertFalse(upstream.hasNext()); return null; }, TestUtils.testEnvBuilder().buildForTrainer() ); dataset.compute(data -> {}); }