com.flipkart.fdp.ml.adapter.VectorAssemblerModelAdapter Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of adapters Show documentation

There is a newer version: 0.1.2

package com.flipkart.fdp.ml.adapter;

import com.flipkart.fdp.ml.modelinfo.VectorAssemblerModelInfo;
import org.apache.spark.ml.feature.VectorAssembler;
import org.apache.spark.sql.DataFrame;

import java.util.Arrays;
import java.util.LinkedHashSet;

/**
 * Transforms Spark's {@link VectorAssembler} in MlLib to  {@link com.flipkart.fdp.ml.modelinfo.VectorAssemblerModelInfo} object
 * that can be exported through {@link com.flipkart.fdp.ml.export.ModelExporter}

 * Created by rohan.shetty on 28/03/16.
 */
public class VectorAssemblerModelAdapter extends AbstractModelInfoAdapter {

    @Override
    VectorAssemblerModelInfo getModelInfo(VectorAssembler from, DataFrame df) {
        VectorAssemblerModelInfo vectorAssemblerModelInfo = new VectorAssemblerModelInfo();
        vectorAssemblerModelInfo.setInputKeys(new LinkedHashSet<>(Arrays.asList(from.getInputCols())));
        vectorAssemblerModelInfo.setOutputKey(from.getOutputCol());
        return vectorAssemblerModelInfo;
    }

    @Override
    public Class getSource() {
        return VectorAssembler.class;
    }

    @Override
    public Class getTarget() {
        return VectorAssemblerModelInfo.class;
    }
}