Search in sources :

Example 1 with RFormulaModel

use of org.apache.spark.ml.feature.RFormulaModel in project jpmml-sparkml by jpmml.

the class RFormulaModelConverter method registerFeatures.

@Override
public void registerFeatures(SparkMLEncoder encoder) {
    RFormulaModel transformer = getTransformer();
    ResolvedRFormula resolvedFormula = transformer.resolvedFormula();
    String targetCol = resolvedFormula.label();
    String labelCol = transformer.getLabelCol();
    if (!(targetCol).equals(labelCol)) {
        List<Feature> features = encoder.getFeatures(targetCol);
        encoder.putFeatures(labelCol, features);
    }
    PipelineModel pipelineModel = transformer.pipelineModel();
    Transformer[] stages = pipelineModel.stages();
    for (Transformer stage : stages) {
        FeatureConverter<?> featureConverter = ConverterUtil.createFeatureConverter(stage);
        featureConverter.registerFeatures(encoder);
    }
}
Also used : Transformer(org.apache.spark.ml.Transformer) ResolvedRFormula(org.apache.spark.ml.feature.ResolvedRFormula) RFormulaModel(org.apache.spark.ml.feature.RFormulaModel) Feature(org.jpmml.converter.Feature) PipelineModel(org.apache.spark.ml.PipelineModel)

Aggregations

PipelineModel (org.apache.spark.ml.PipelineModel)1 Transformer (org.apache.spark.ml.Transformer)1 RFormulaModel (org.apache.spark.ml.feature.RFormulaModel)1 ResolvedRFormula (org.apache.spark.ml.feature.ResolvedRFormula)1 Feature (org.jpmml.converter.Feature)1