diff --git a/pmml-xgboost/src/main/java/org/jpmml/xgboost/Learner.java b/pmml-xgboost/src/main/java/org/jpmml/xgboost/Learner.java index c5c07b8..541f78c 100644 --- 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 { 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 splitTypes = Learner.this.gbtree.getSplitType(splitIndex); if(splitTypes.size() == 0){ - return Node.SPLIT_NUMERICAL; + return null; } else if(splitTypes.size() == 1){