[ML] fixing minor text categorization aggregation bug (#78629)

elastic · Oct 4, 2021 · 81aa483 · 81aa483
1 parent f6fbeb8
commit 81aa483
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/TreeNode.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/TreeNode.java
@@ -252,7 +252,7 @@ void collapseTinyChildren() {
                 return;
             }
             Optional<TreeNode> maybeWildChild = getChild(WILD_CARD_ID).or(() -> {
-                if ((double) smallestChild.peek().count / this.getCount() <= 1.0 / maxChildren) {
+                if (smallestChild.size() > 0 && (double) smallestChild.peek().count / this.getCount() <= 1.0 / maxChildren) {
                     TreeNode tinyChild = children.remove(smallestChild.poll().tokenId);
                     return Optional.of(addChild(WILD_CARD_ID, tinyChild));
                 }
@@ -267,7 +267,7 @@ void collapseTinyChildren() {
                         smallestChild.add(tinyNode);
                         break;
                     } else {
-                        wildChild.mergeWith(children.remove(tinyNode.count));
+                        wildChild.mergeWith(children.remove(tinyNode.tokenId));
                     }
                 }
             }