com.flipkart.fdp.ml.adapter.VectorAssemblerModelAdapter Maven / Gradle / Ivy
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;
import java.util.Set;
/**
* 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())));
Set outputKeys = new LinkedHashSet();
outputKeys.add(from.getOutputCol());
vectorAssemblerModelInfo.setOutputKeys(outputKeys);
return vectorAssemblerModelInfo;
}
@Override
public Class getSource() {
return VectorAssembler.class;
}
@Override
public Class getTarget() {
return VectorAssemblerModelInfo.class;
}
}