Added FeatureMap$Entry$Type enum

jpmml · May 23, 2020 · 04222a3 · 04222a3
1 parent 8bcb79a
commit 04222a3
Showing 1 changed file with 33 additions and 10 deletions.
diff --git a/src/main/java/org/jpmml/xgboost/FeatureMap.java b/src/main/java/org/jpmml/xgboost/FeatureMap.java
@@ -61,20 +61,20 @@ public List<Feature> 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);
+				}
+			}
+		}
 	}
 }