Added activation condition

jpmml · Nov 14, 2021 · 2ac12bb · 2ac12bb
1 parent 5225310
commit 2ac12bb
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/src/main/java/org/jpmml/translator/mining/TreeModelBoosterTranslator.java b/src/main/java/org/jpmml/translator/mining/TreeModelBoosterTranslator.java
@@ -21,6 +21,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import com.sun.codemodel.JBlock;
 import com.sun.codemodel.JExpr;
@@ -125,6 +126,23 @@ public TreeModelBoosterTranslator(PMML pmml, MiningModel miningModel){
 					throw new UnsupportedAttributeException(treeModel, missingValueStrategy);
 			}
 		}
+
+		AtomicInteger nodeCount = new AtomicInteger(0);
+
+		Visitor nodeCounter = new AbstractVisitor(){
+
+			@Override
+			public VisitorAction visit(Node node){
+				nodeCount.incrementAndGet();
+
+				return super.visit(node);
+			}
+		};
+		nodeCounter.applyTo(segmentation);
+
+		if(nodeCount.get() > TreeModelBoosterTranslator.NODE_COUNT_LIMIT){
+			throw new UnsupportedElementException(segmentation);
+		}
 	}
 
 	@Override
@@ -387,4 +405,6 @@ private void addNode(Node node){
 
 		return treeModel;
 	}
+
+	public static final int NODE_COUNT_LIMIT = Integer.getInteger(TreeModelBoosterTranslator.class.getName() + "#" + "NODE_COUNT_LIMIT", 1000);
 }