Improved support for unused features

jpmml · Dec 18, 2023 · 3f9bec7 · 3f9bec7
1 parent 8f5f701
commit 3f9bec7
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/pmml-xgboost/src/main/java/org/jpmml/xgboost/Learner.java b/pmml-xgboost/src/main/java/org/jpmml/xgboost/Learner.java
@@ -674,7 +674,11 @@ private class FeatureTransformer implements Function<Feature, Feature> {
 		public Feature apply(Feature feature){
 			int splitIndex = getSplitIndex(feature);
 
-			int splitType = getSplitType(splitIndex);
+			Integer splitType = getSplitType(splitIndex);
+			if(splitType == null){
+				return feature;
+			}
+
 			switch(splitType){
 				case Node.SPLIT_NUMERICAL:
 					return transformNumerical(feature);
@@ -685,11 +689,11 @@ public Feature apply(Feature feature){
 			}
 		}
 
-		private int getSplitType(int splitIndex){
+		private Integer getSplitType(int splitIndex){
 			Set<Integer> splitTypes = Learner.this.gbtree.getSplitType(splitIndex);
 
 			if(splitTypes.size() == 0){
-				return Node.SPLIT_NUMERICAL;
+				return null;
 			} else
 
 			if(splitTypes.size() == 1){