ImputerTrainer implements PreprocessingTrainer<K, V> { @Override public ImputerPreprocessor<K, V> fit(LearningEnvironmentBuilder envBuilder, DatasetBuilder<K, V> datasetBuilder, Preprocessor<K, V> basePreprocessor) { PartitionContextBuilder<K, V, EmptyContext> builder = (env, upstream, upstreamSize) -> new EmptyContext(); try (Dataset<EmptyContext, ImputerPartitionData> dataset = datasetBuilder.build( envBuilder, builder, (env, upstream, upstreamSize, ctx) -> { double[] sums = null; int[] counts = null; Map<Double, Integer>[] valuesByFreq = null; while (upstream.hasNext()) { UpstreamEntry<K, V> entity = upstream.next(); LabeledVector row = basePreprocessor.apply(entity.getKey(), entity.getValue()); switch (imputingStgy) { case MEAN: sums = calculateTheSums(row, sums); counts = calculateTheCounts(row, counts); break; case MOST_FREQUENT: valuesByFreq = calculateFrequencies(row, valuesByFreq); break; default: throw new UnsupportedOperationException("The chosen strategy is not supported"); } } ImputerPartitionData partData; switch (imputingStgy) { case MEAN: partData = new ImputerPartitionData().withSums(sums).withCounts(counts); break; case MOST_FREQUENT: partData = new ImputerPartitionData().withValuesByFrequency(valuesByFreq); break; default: throw new UnsupportedOperationException("The chosen strategy is not supported"); } return partData; }, learningEnvironment(basePreprocessor) )) { Vector imputingValues; switch (imputingStgy) { case MEAN: imputingValues = VectorUtils.of(calculateImputingValuesBySumsAndCounts(dataset)); break; case MOST_FREQUENT: imputingValues = VectorUtils.of(calculateImputingValuesByFrequencies(dataset)); break; default: throw new UnsupportedOperationException("The chosen strategy is not supported"); } return new ImputerPreprocessor<>(imputingValues, basePreprocessor); } catch (Exception e) { throw new RuntimeException(e); } }  @Override ImputerPreprocessor<K, V> fit(LearningEnvironmentBuilder envBuilder, DatasetBuilder<K, V> datasetBuilder,
                                                   Preprocessor<K, V> basePreprocessor); ImputerTrainer<K, V> withImputingStrategy(ImputingStrategy imputingStgy);  }
@Test public void testFit() { Map<Integer, Vector> data = new HashMap<>(); data.put(1, VectorUtils.of(1, 2, Double.NaN)); data.put(2, VectorUtils.of(1, Double.NaN, 22)); data.put(3, VectorUtils.of(Double.NaN, 10, 100)); data.put(4, VectorUtils.of(0, 2, 100)); DatasetBuilder<Integer, Vector> datasetBuilder = new LocalDatasetBuilder<>(data, parts); final Vectorizer<Integer, Vector, Integer, Double> vectorizer = new DummyVectorizer<>(0, 1, 2); ImputerTrainer<Integer, Vector> imputerTrainer = new ImputerTrainer<Integer, Vector>() .withImputingStrategy(ImputingStrategy.MOST_FREQUENT); ImputerPreprocessor<Integer, Vector> preprocessor = imputerTrainer.fit( TestUtils.testEnvBuilder(), datasetBuilder, vectorizer ); assertArrayEquals(new double[] {1, 0, 100}, preprocessor.apply(5, VectorUtils.of(Double.NaN, 0, Double.NaN)).features().asArray(), 1e-8); }