Small refactoring in tree building

Co-authored-by: Veselin Nikolov <[email protected]>

Author: IoannisPanagiotas

algo/src/main/java/org/neo4j/gds/hdbscan/KdNode.java

@@ -50,7 +50,6 @@ private KdNode(long id, long start, long end, AABB aabb, boolean isLeaf) {
         this.isLeaf = isLeaf;
     }
 
-
     void leftChild(KdNode child){
         this.leftChild = child;
     }

algo/src/main/java/org/neo4j/gds/hdbscan/KdTreeBuilder.java

@@ -67,6 +67,8 @@ public KdTree build(){
             nodePropertyValues.nodeCount(),
             leafSize,
             nodeIndex,
+            null,
+            false,
             progressTracker
         );
 

algo/src/main/java/org/neo4j/gds/hdbscan/KdTreeNodeBuilderTask.java

@@ -37,6 +37,8 @@ class KdTreeNodeBuilderTask implements Runnable {
     private final AtomicInteger  nodeIndex;
     private final ProgressTracker progressTracker;
     private  final KDNodeSupport kdNodeSupport;
+    private final KdNode parent;
+    private final boolean amLeftChild;
 
 
     KdTreeNodeBuilderTask(
@@ -47,6 +49,8 @@ class KdTreeNodeBuilderTask implements Runnable {
         long end,
         long maxLeafSize,
         AtomicInteger nodeIndex,
+        KdNode parent,
+        boolean amLeftChild,
         ProgressTracker progressTracker
     ) {
         this.ids = ids;
@@ -57,6 +61,8 @@ class KdTreeNodeBuilderTask implements Runnable {
         this.nodeIndex = nodeIndex;
         this.progressTracker = progressTracker;
         this.kdNodeSupport = kdNodeSupport;
+        this.parent = parent;
+        this.amLeftChild = amLeftChild;
     }
 
     @Override
@@ -65,36 +71,45 @@ public void run() {
         var aabb = kdNodeSupport.create(start, end);
         var treeNodeId = nodeIndex.getAndIncrement();
         if (nodeSize <= maxLeafSize) {
-            kdNode = KdNode.createLeaf(treeNodeId, start, end, aabb);
+            createNode(true,treeNodeId,aabb);
             progressTracker.logProgress(nodeSize);
         } else {
 
             int indexToSplit = aabb.mostSpreadDimension(); //step. 1: find the index to  dimension split
             long median = findMedianAndSplit(indexToSplit);  //step.2  modify array so that everything < is before median and everything >= after
 
-            kdNode = KdNode.createSplitNode(
-                treeNodeId,
-                start,
-                end,
-                aabb
-            );
+            createNode(false,treeNodeId,aabb);
+
             //TODO: step.4 add these builder tasks into a fork-join
             var leftChildBuilder = new KdTreeNodeBuilderTask(
-                kdNodeSupport, nodePropertyValues, ids, start, median, maxLeafSize, nodeIndex, progressTracker
+                kdNodeSupport,
+                nodePropertyValues,
+                ids,
+                start,
+                median,
+                maxLeafSize,
+                nodeIndex,
+                kdNode,
+                true,
+                progressTracker
             );
             leftChildBuilder.run();
 
             var rightChildBuilder = new KdTreeNodeBuilderTask(
-                kdNodeSupport, nodePropertyValues, ids, median, end, maxLeafSize, nodeIndex, progressTracker
+                kdNodeSupport,
+                nodePropertyValues,
+                ids,
+                median,
+                end,
+                maxLeafSize,
+                nodeIndex,
+                kdNode,
+                false,
+                progressTracker
             );
 
             rightChildBuilder.run();
 
-            var leftChild = leftChildBuilder.kdNode();
-            var rightChild = rightChildBuilder.kdNode();
-
-            kdNode.leftChild(leftChild);
-            kdNode.rightChild(rightChild);
         }
     }
 
@@ -151,5 +166,29 @@ private long partition(LongToDoubleFunction valueAt, long l, long r) {
         return pIndex;
     }
 
+    private void createNode(boolean isLeaf, long treeNodeId, AABB aabb){
+        if (isLeaf) {
+            kdNode = KdNode.createLeaf(
+                treeNodeId,
+                start,
+                end, aabb
+            );
+        }else{
+            kdNode = KdNode.createSplitNode(
+                treeNodeId,
+                start,
+                end,
+                aabb
+            );
+        }
+        if (parent!=null){
+            if (amLeftChild){
+                parent.leftChild(kdNode);
+            }else{
+                parent.rightChild(kdNode);
+            }
+        }
+    }
+
 
 }

algo/src/test/java/org/neo4j/gds/hdbscan/KdTreeNodeBuilderTaskTest.java

@@ -61,6 +61,8 @@ public long nodeCount() {
             10,
             1,
             new AtomicInteger(0),
+            null,
+            false,
             ProgressTracker.NULL_TRACKER
         );
 
@@ -98,6 +100,8 @@ public long nodeCount() {
             8,
             1,
             new AtomicInteger(),
+            null,
+            false,
             ProgressTracker.NULL_TRACKER
         );
 
@@ -132,6 +136,8 @@ public long nodeCount() {
             3,
             3,
             new AtomicInteger(),
+            null,
+            false,
             ProgressTracker.NULL_TRACKER
         );
 
@@ -163,14 +169,16 @@ public long nodeCount() {
 
         var aabbCreator = new DoubleKDNodeSupport(nodePropertyValues,ids,1);
 
-        var nodeBuilder =new KdTreeNodeBuilderTask(
+        var nodeBuilder = new KdTreeNodeBuilderTask(
             aabbCreator,
             nodePropertyValues,
             ids,
             0,
             3,
             2,
             new AtomicInteger(0),
+            null,
+            false,
             ProgressTracker.NULL_TRACKER
         );