streamlined vector value extractor

Author: sonalgoyal

common/core/src/test/java/zingg/common/core/executor/MatchThresholdBasedExecutorTester.java

@@ -37,7 +37,7 @@ protected void runUntilThreshold() {
             LOG.info("total number of matches discovered, " + matchCount);
             LOG.info("total number of non-matches discovered, " + notAMatchCount);
         } catch (Exception | ZinggClientException exception) {
-            throw new ZinggException("Exception occurred while running threshold based runner, " + exception.getMessage());
+            throw new ZinggException("Exception occurred while running threshold based runner, " + exception.getMessage(), exception);
         }
     }
 

spark/core/src/main/java/zingg/spark/core/model/SparkMLPipeline.java

@@ -57,18 +57,25 @@ public SparkMLPipeline(Map<FieldDefinition, Feature<DataType>> featurers, ModelC
 	public void buildPipeline(Map<FieldDefinition, Feature<DataType>> featurers, ModelColumnHelper columnHelper) {
 		pipelineStage = new ArrayList<>();
 
-		featureCreators = new SparkFeatureCreators(featurers, columnHelper);
+		featureCreators = getFeatureCreators(featurers);
 		pipelineStage.addAll(featureCreators.getTransformers());
 
 		pipelineStage.add(getAssembler());
 		pipelineStage.add(getPolyExpansion());
 		pipelineStage.add(getLR());
 
-		vve = new VectorValueExtractor(ColName.PROBABILITY_COL, ColName.SCORE_COL);
-		columnHelper.getColumnsAdded().add(ColName.PROBABILITY_COL);
-		columnHelper.getColumnsAdded().add(ColName.RAW_PREDICTION);
+		//vve is not used in all cases, but since we dont have a different flow for prediction
+		//creating it here so it gets registered
+		vve = getVVE();
+
+		
 	}
 
+	protected SparkFeatureCreators getFeatureCreators(Map<FieldDefinition, Feature<DataType>> featurers){
+		featureCreators = new SparkFeatureCreators(featurers, columnHelper);
+		return featureCreators;
+	} 
+
 	protected VectorAssembler getAssembler() {
 		VectorAssembler assembler = new VectorAssembler();
 		assembler.setInputCols(columnHelper.getColumnsAdded().toArray(new String[0]))
@@ -123,11 +130,18 @@ public ZFrame<Dataset<Row>, Row, Column> fit(ZFrame<Dataset<Row>, Row, Column> i
 	public ZFrame<Dataset<Row>, Row, Column> predict(ZFrame<Dataset<Row>, Row, Column> data) {
 		LOG.info("threshold while predicting is " + lr.getThreshold());
 		Dataset<Row> predictWithFeatures = transformer.transform(data.df());
-		predictWithFeatures = vve.transform(predictWithFeatures);
+		predictWithFeatures = getVVE().transform(predictWithFeatures);
 		LOG.debug("Return schema is " + predictWithFeatures.schema());
 		return new SparkFrame(predictWithFeatures);
 	}
 
+	public VectorValueExtractor getVVE() {
+		vve = new VectorValueExtractor(ColName.PROBABILITY_COL, ColName.SCORE_COL);
+		columnHelper.getColumnsAdded().add(ColName.PROBABILITY_COL);
+		columnHelper.getColumnsAdded().add(ColName.RAW_PREDICTION);
+		return vve;
+	}
+
 	public void register(SparkSession session) {
 		featureCreators.register(session);
 		vve.register(session);
@@ -141,7 +155,5 @@ public void save(String path) throws IOException {
 		((CrossValidatorModel) transformer).write().overwrite().save(path);
 	}
 
-	public List<SparkTransformer> getFeatureCreators() {
-		return featureCreators.getTransformers();
-	}
+	
 }