Skip to content

Commit 89365ce

Browse files
Small refactoring in tree building
Co-authored-by: Veselin Nikolov <[email protected]>
1 parent 79b1711 commit 89365ce

File tree

4 files changed

+64
-16
lines changed

4 files changed

+64
-16
lines changed

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,6 @@ private KdNode(long id, long start, long end, AABB aabb, boolean isLeaf) {
5050
this.isLeaf = isLeaf;
5151
}
5252

53-
5453
void leftChild(KdNode child){
5554
this.leftChild = child;
5655
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,8 @@ public KdTree build(){
6767
nodePropertyValues.nodeCount(),
6868
leafSize,
6969
nodeIndex,
70+
null,
71+
false,
7072
progressTracker
7173
);
7274

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

Lines changed: 53 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ class KdTreeNodeBuilderTask implements Runnable {
3737
private final AtomicInteger nodeIndex;
3838
private final ProgressTracker progressTracker;
3939
private final KDNodeSupport kdNodeSupport;
40+
private final KdNode parent;
41+
private final boolean amLeftChild;
4042

4143

4244
KdTreeNodeBuilderTask(
@@ -47,6 +49,8 @@ class KdTreeNodeBuilderTask implements Runnable {
4749
long end,
4850
long maxLeafSize,
4951
AtomicInteger nodeIndex,
52+
KdNode parent,
53+
boolean amLeftChild,
5054
ProgressTracker progressTracker
5155
) {
5256
this.ids = ids;
@@ -57,6 +61,8 @@ class KdTreeNodeBuilderTask implements Runnable {
5761
this.nodeIndex = nodeIndex;
5862
this.progressTracker = progressTracker;
5963
this.kdNodeSupport = kdNodeSupport;
64+
this.parent = parent;
65+
this.amLeftChild = amLeftChild;
6066
}
6167

6268
@Override
@@ -65,36 +71,45 @@ public void run() {
6571
var aabb = kdNodeSupport.create(start, end);
6672
var treeNodeId = nodeIndex.getAndIncrement();
6773
if (nodeSize <= maxLeafSize) {
68-
kdNode = KdNode.createLeaf(treeNodeId, start, end, aabb);
74+
createNode(true,treeNodeId,aabb);
6975
progressTracker.logProgress(nodeSize);
7076
} else {
7177

7278
int indexToSplit = aabb.mostSpreadDimension(); //step. 1: find the index to dimension split
7379
long median = findMedianAndSplit(indexToSplit); //step.2 modify array so that everything < is before median and everything >= after
7480

75-
kdNode = KdNode.createSplitNode(
76-
treeNodeId,
77-
start,
78-
end,
79-
aabb
80-
);
81+
createNode(false,treeNodeId,aabb);
82+
8183
//TODO: step.4 add these builder tasks into a fork-join
8284
var leftChildBuilder = new KdTreeNodeBuilderTask(
83-
kdNodeSupport, nodePropertyValues, ids, start, median, maxLeafSize, nodeIndex, progressTracker
85+
kdNodeSupport,
86+
nodePropertyValues,
87+
ids,
88+
start,
89+
median,
90+
maxLeafSize,
91+
nodeIndex,
92+
kdNode,
93+
true,
94+
progressTracker
8495
);
8596
leftChildBuilder.run();
8697

8798
var rightChildBuilder = new KdTreeNodeBuilderTask(
88-
kdNodeSupport, nodePropertyValues, ids, median, end, maxLeafSize, nodeIndex, progressTracker
99+
kdNodeSupport,
100+
nodePropertyValues,
101+
ids,
102+
median,
103+
end,
104+
maxLeafSize,
105+
nodeIndex,
106+
kdNode,
107+
false,
108+
progressTracker
89109
);
90110

91111
rightChildBuilder.run();
92112

93-
var leftChild = leftChildBuilder.kdNode();
94-
var rightChild = rightChildBuilder.kdNode();
95-
96-
kdNode.leftChild(leftChild);
97-
kdNode.rightChild(rightChild);
98113
}
99114
}
100115

@@ -151,5 +166,29 @@ private long partition(LongToDoubleFunction valueAt, long l, long r) {
151166
return pIndex;
152167
}
153168

169+
private void createNode(boolean isLeaf, long treeNodeId, AABB aabb){
170+
if (isLeaf) {
171+
kdNode = KdNode.createLeaf(
172+
treeNodeId,
173+
start,
174+
end, aabb
175+
);
176+
}else{
177+
kdNode = KdNode.createSplitNode(
178+
treeNodeId,
179+
start,
180+
end,
181+
aabb
182+
);
183+
}
184+
if (parent!=null){
185+
if (amLeftChild){
186+
parent.leftChild(kdNode);
187+
}else{
188+
parent.rightChild(kdNode);
189+
}
190+
}
191+
}
192+
154193

155194
}

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,8 @@ public long nodeCount() {
6161
10,
6262
1,
6363
new AtomicInteger(0),
64+
null,
65+
false,
6466
ProgressTracker.NULL_TRACKER
6567
);
6668

@@ -98,6 +100,8 @@ public long nodeCount() {
98100
8,
99101
1,
100102
new AtomicInteger(),
103+
null,
104+
false,
101105
ProgressTracker.NULL_TRACKER
102106
);
103107

@@ -132,6 +136,8 @@ public long nodeCount() {
132136
3,
133137
3,
134138
new AtomicInteger(),
139+
null,
140+
false,
135141
ProgressTracker.NULL_TRACKER
136142
);
137143

@@ -163,14 +169,16 @@ public long nodeCount() {
163169

164170
var aabbCreator = new DoubleKDNodeSupport(nodePropertyValues,ids,1);
165171

166-
var nodeBuilder =new KdTreeNodeBuilderTask(
172+
var nodeBuilder = new KdTreeNodeBuilderTask(
167173
aabbCreator,
168174
nodePropertyValues,
169175
ids,
170176
0,
171177
3,
172178
2,
173179
new AtomicInteger(0),
180+
null,
181+
false,
174182
ProgressTracker.NULL_TRACKER
175183
);
176184

0 commit comments

Comments
 (0)