Pipeline extends EstimatorBase<Pipeline, PipelineModel> { @Override public PipelineModel fit(BatchOperator input) { int lastEstimatorIdx = getIndexOfLastEstimator(); TransformerBase[] transformers = new TransformerBase[stages.size()]; for (int i = 0; i < stages.size(); i++) { PipelineStageBase stage = stages.get(i); if (i <= lastEstimatorIdx) { if (stage instanceof EstimatorBase) { transformers[i] = ((EstimatorBase) stage).fit(input); } else if (stage instanceof TransformerBase) { transformers[i] = (TransformerBase) stage; } if (i < lastEstimatorIdx) { Boolean lazyPrintTransformDataEnabled = (Boolean) transformers[i].get(LAZY_PRINT_TRANSFORM_DATA_ENABLED); Boolean lazyPrintTransformStatEnabled = (Boolean) transformers[i].get(LAZY_PRINT_TRANSFORM_STAT_ENABLED); transformers[i].set(LAZY_PRINT_TRANSFORM_DATA_ENABLED, false); transformers[i].set(LAZY_PRINT_TRANSFORM_STAT_ENABLED, false); input = transformers[i].transform(input); transformers[i].set(LAZY_PRINT_TRANSFORM_DATA_ENABLED, lazyPrintTransformDataEnabled); transformers[i].set(LAZY_PRINT_TRANSFORM_STAT_ENABLED, lazyPrintTransformStatEnabled); } } else { transformers[i] = (TransformerBase) stage; } } return new PipelineModel(transformers).setMLEnvironmentId(input.getMLEnvironmentId()); } Pipeline();  Pipeline(Params params);  Pipeline(PipelineStageBase<?>... stages); @Override Pipeline clone(); Pipeline add(PipelineStageBase stage); Pipeline add(int index, PipelineStageBase stage); Pipeline remove(int index); PipelineStageBase get(int index); int size(); @Override PipelineModel fit(BatchOperator input); @Override PipelineModel fit(StreamOperator input);  }
@Test public void testFit() { BatchOperator data = new MemSourceBatchOp(new Object[]{1}, "colName"); TransformerBase stage1 = mockTransformer("stage1"); TransformerBase stage2 = mockTransformer("stage2"); Pair<EstimatorBase, ModelBase> stage3 = mockEstimator("stage3"); TransformerBase stage4 = mockTransformer("stage4"); Pair<EstimatorBase, ModelBase> stage5 = mockEstimator("stage5"); TransformerBase stage6 = mockTransformer("stage6"); Pipeline pipe = new Pipeline().add(stage1).add(stage2).add(stage3.getLeft()) .add(stage4).add(stage5.getLeft()).add(stage6); pipe.fit(data); verify(stage1, times(1)).transform(any(BatchOperator.class)); verify(stage2, times(1)).transform(any(BatchOperator.class)); verify(stage3.getLeft(), times(1)).fit(any(BatchOperator.class)); verify(stage3.getRight(), times(1)).transform(any(BatchOperator.class)); verify(stage4, times(1)).transform(any(BatchOperator.class)); verify(stage5.getLeft(), times(1)).fit(any(BatchOperator.class)); verify(stage5.getRight(), never()).transform(any(BatchOperator.class)); verify(stage6, never()).transform(any(BatchOperator.class)); }