diff --git a/src/main/java/org/jpmml/xgboost/FeatureMap.java b/src/main/java/org/jpmml/xgboost/FeatureMap.java index 934f29a..870434b 100644 --- a/src/main/java/org/jpmml/xgboost/FeatureMap.java +++ b/src/main/java/org/jpmml/xgboost/FeatureMap.java @@ -61,20 +61,20 @@ public List encodeFeatures(PMMLEncoder encoder){ DataField dataField = encoder.getDataField(name); if(dataField == null){ - String type = entry.getType(); + Entry.Type type = entry.getType(); switch(type){ - case "i": + case BINARY_INDICATOR: dataField = encoder.createDataField(name, OpType.CATEGORICAL, DataType.STRING); break; - case "q": + case FLOAT: dataField = encoder.createDataField(name, OpType.CONTINUOUS, DataType.FLOAT); break; - case "int": + case INTEGER: dataField = encoder.createDataField(name, OpType.CONTINUOUS, DataType.INTEGER); break; default: - throw new IllegalArgumentException(type); + throw new IllegalArgumentException(String.valueOf(type)); } } @@ -127,7 +127,7 @@ public void addEntry(String name, String type){ name = name.substring(0, equals); } - Entry entry = new Entry(name, value, type); + Entry entry = new Entry(name, value, Entry.Type.fromString(type)); addEntry(entry); } @@ -177,10 +177,10 @@ public class Entry { private String value = null; - private String type = null; + private Type type = null; - public Entry(String name, String value, String type){ + public Entry(String name, String value, Type type){ setName(name); setValue(value); setType(type); @@ -202,12 +202,35 @@ private void setValue(String value){ this.value = value; } - public String getType(){ + public Type getType(){ return this.type; } - private void setType(String type){ + private void setType(Type type){ this.type = Objects.requireNonNull(type); } + + static + public enum Type { + BINARY_INDICATOR, + FLOAT, + INTEGER, + ; + + static + public Type fromString(String string){ + + switch(string){ + case "i": + return Type.BINARY_INDICATOR; + case "q": + return Type.FLOAT; + case "int": + return Type.INTEGER; + default: + throw new IllegalArgumentException(string); + } + } + } } } \ No newline at end of file