Search in sources :

Example 1 with OneHotEncoder

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

the class OneHotEncoderConverter method encodeFeatures.

@Override
public List<Feature> encodeFeatures(SparkMLEncoder encoder) {
    OneHotEncoder transformer = getTransformer();
    boolean dropLast = true;
    Option<Object> dropLastOption = transformer.get(transformer.dropLast());
    if (dropLastOption.isDefined()) {
        dropLast = (Boolean) dropLastOption.get();
    }
    CategoricalFeature categoricalFeature = (CategoricalFeature) encoder.getOnlyFeature(transformer.getInputCol());
    List<String> values = categoricalFeature.getValues();
    if (dropLast) {
        values = values.subList(0, values.size() - 1);
    }
    List<Feature> result = new ArrayList<>();
    for (String value : values) {
        result.add(new BinaryFeature(encoder, categoricalFeature.getName(), DataType.STRING, value));
    }
    return result;
}
Also used : OneHotEncoder(org.apache.spark.ml.feature.OneHotEncoder) ArrayList(java.util.ArrayList) BinaryFeature(org.jpmml.converter.BinaryFeature) Feature(org.jpmml.converter.Feature) CategoricalFeature(org.jpmml.converter.CategoricalFeature) BinaryFeature(org.jpmml.converter.BinaryFeature) CategoricalFeature(org.jpmml.converter.CategoricalFeature)

Aggregations

ArrayList (java.util.ArrayList)1 OneHotEncoder (org.apache.spark.ml.feature.OneHotEncoder)1 BinaryFeature (org.jpmml.converter.BinaryFeature)1 CategoricalFeature (org.jpmml.converter.CategoricalFeature)1 Feature (org.jpmml.converter.Feature)1