Skip to content
This repository was archived by the owner on Apr 22, 2020. It is now read-only.

Commit 4dcc706

Browse files
committed
adding tests for similarity exporter
1 parent 7914cd5 commit 4dcc706

File tree

5 files changed

+254
-135
lines changed

5 files changed

+254
-135
lines changed

algo/src/main/java/org/neo4j/graphalgo/similarity/ParallelSimilarityExporter.java

+4-3
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@
4646
import java.util.stream.Collectors;
4747
import java.util.stream.Stream;
4848

49-
public class ParallelSimilarityExporter extends StatementApi {
49+
public class ParallelSimilarityExporter extends StatementApi implements SimilarityExporter {
5050

5151
private final Log log;
5252
private final int propertyId;
@@ -64,7 +64,7 @@ public ParallelSimilarityExporter(GraphDatabaseAPI api,
6464
this.nodeCount = nodeCount;
6565
}
6666

67-
public void export(Stream<SimilarityResult> similarityPairs, long batchSize) {
67+
public int export(Stream<SimilarityResult> similarityPairs, long batchSize) {
6868
IdMap idMap = new IdMap(this.nodeCount);
6969
AdjacencyMatrix adjacencyMatrix = new AdjacencyMatrix(this.nodeCount, false, AllocationTracker.EMPTY);
7070
WeightMap weightMap = new WeightMap(nodeCount, 0, propertyId);
@@ -137,6 +137,7 @@ public void export(Stream<SimilarityResult> similarityPairs, long batchSize) {
137137

138138
int outQueueBatches = writeSequential(outQueue.stream().flatMap(Collection::stream), batchSize);
139139
log.info("ParallelSimilarityExporter: Batch Size: %d, Batches written - in parallel: %d, sequentially: %d", batchSize, inQueueBatches, outQueueBatches);
140+
return inQueueBatches + outQueueBatches;
140141
}
141142

142143
private static <T> void put(BlockingQueue<T> queue, T items) {
@@ -213,7 +214,7 @@ private int writeSequential(Stream<SimilarityResult> similarityPairs, long batch
213214
do {
214215
List<SimilarityResult> batch = take(iterator, Math.toIntExact(batchSize));
215216
export(batch);
216-
if(batch.size() > 0) {
217+
if (batch.size() > 0) {
217218
counter[0]++;
218219
}
219220
} while (iterator.hasNext());
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
/**
2+
* Copyright (c) 2017 "Neo4j, Inc." <http://neo4j.com>
3+
* <p>
4+
* This file is part of Neo4j Graph Algorithms <http://github.com/neo4j-contrib/neo4j-graph-algorithms>.
5+
* <p>
6+
* Neo4j Graph Algorithms is free software: you can redistribute it and/or modify
7+
* it under the terms of the GNU General Public License as published by
8+
* the Free Software Foundation, either version 3 of the License, or
9+
* (at your option) any later version.
10+
* <p>
11+
* This program is distributed in the hope that it will be useful,
12+
* but WITHOUT ANY WARRANTY; without even the implied warranty of
13+
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14+
* GNU General Public License for more details.
15+
* <p>
16+
* You should have received a copy of the GNU General Public License
17+
* along with this program. If not, see <http://www.gnu.org/licenses/>.
18+
*/
19+
package org.neo4j.graphalgo.similarity;
20+
21+
import org.neo4j.graphalgo.core.utils.ExceptionUtil;
22+
import org.neo4j.graphalgo.core.utils.StatementApi;
23+
import org.neo4j.internal.kernel.api.exceptions.EntityNotFoundException;
24+
import org.neo4j.internal.kernel.api.exceptions.InvalidTransactionTypeKernelException;
25+
import org.neo4j.internal.kernel.api.exceptions.KernelException;
26+
import org.neo4j.internal.kernel.api.exceptions.explicitindex.AutoIndexingKernelException;
27+
import org.neo4j.kernel.api.KernelTransaction;
28+
import org.neo4j.kernel.internal.GraphDatabaseAPI;
29+
import org.neo4j.logging.Log;
30+
import org.neo4j.values.storable.Values;
31+
32+
import java.util.ArrayList;
33+
import java.util.Iterator;
34+
import java.util.List;
35+
import java.util.stream.Stream;
36+
37+
public class SequentialSimilarityExporter extends StatementApi implements SimilarityExporter {
38+
39+
private final Log log;
40+
private final int propertyId;
41+
private final int relationshipTypeId;
42+
43+
public SequentialSimilarityExporter(GraphDatabaseAPI api,
44+
Log log, String relationshipType,
45+
String propertyName) {
46+
super(api);
47+
this.log = log;
48+
propertyId = getOrCreatePropertyId(propertyName);
49+
relationshipTypeId = getOrCreateRelationshipId(relationshipType);
50+
}
51+
52+
public int export(Stream<SimilarityResult> similarityPairs, long batchSize) {
53+
int batches = writeSequential(similarityPairs, batchSize);
54+
log.info("SequentialSimilarityExporter: Batch Size: %d, Batches written - sequentially: %d", batchSize, batches);
55+
return batches;
56+
}
57+
58+
private void export(SimilarityResult similarityResult) {
59+
applyInTransaction(statement -> {
60+
try {
61+
createRelationship(similarityResult, statement);
62+
} catch (KernelException e) {
63+
ExceptionUtil.throwKernelException(e);
64+
}
65+
return null;
66+
});
67+
68+
}
69+
70+
private void export(List<SimilarityResult> similarityResults) {
71+
applyInTransaction(statement -> {
72+
for (SimilarityResult similarityResult : similarityResults) {
73+
try {
74+
createRelationship(similarityResult, statement);
75+
} catch (KernelException e) {
76+
ExceptionUtil.throwKernelException(e);
77+
}
78+
}
79+
return null;
80+
});
81+
82+
}
83+
84+
private void createRelationship(SimilarityResult similarityResult, KernelTransaction statement) throws EntityNotFoundException, InvalidTransactionTypeKernelException, AutoIndexingKernelException {
85+
long node1 = similarityResult.item1;
86+
long node2 = similarityResult.item2;
87+
long relationshipId = statement.dataWrite().relationshipCreate(node1, relationshipTypeId, node2);
88+
89+
statement.dataWrite().relationshipSetProperty(
90+
relationshipId, propertyId, Values.doubleValue(similarityResult.similarity));
91+
}
92+
93+
private int getOrCreateRelationshipId(String relationshipType) {
94+
return applyInTransaction(stmt -> stmt
95+
.tokenWrite()
96+
.relationshipTypeGetOrCreateForName(relationshipType));
97+
}
98+
99+
private int getOrCreatePropertyId(String propertyName) {
100+
return applyInTransaction(stmt -> stmt
101+
.tokenWrite()
102+
.propertyKeyGetOrCreateForName(propertyName));
103+
}
104+
105+
private int writeSequential(Stream<SimilarityResult> similarityPairs, long batchSize) {
106+
log.info("SequentialSimilarityExporter: Writing relationships...");
107+
int[] counter = {0};
108+
if (batchSize == 1) {
109+
similarityPairs.forEach(similarityResult -> {
110+
export(similarityResult);
111+
counter[0]++;
112+
});
113+
} else {
114+
Iterator<SimilarityResult> iterator = similarityPairs.iterator();
115+
do {
116+
List<SimilarityResult> batch = take(iterator, Math.toIntExact(batchSize));
117+
export(batch);
118+
if(batch.size() > 0) {
119+
counter[0]++;
120+
}
121+
} while (iterator.hasNext());
122+
}
123+
124+
return counter[0];
125+
}
126+
127+
128+
private static List<SimilarityResult> take(Iterator<SimilarityResult> iterator, int batchSize) {
129+
List<SimilarityResult> result = new ArrayList<>(batchSize);
130+
while (iterator.hasNext() && batchSize-- > 0) {
131+
result.add(iterator.next());
132+
}
133+
return result;
134+
}
135+
136+
137+
}
Original file line numberDiff line numberDiff line change
@@ -1,136 +1,7 @@
1-
/**
2-
* Copyright (c) 2017 "Neo4j, Inc." <http://neo4j.com>
3-
* <p>
4-
* This file is part of Neo4j Graph Algorithms <http://github.com/neo4j-contrib/neo4j-graph-algorithms>.
5-
* <p>
6-
* Neo4j Graph Algorithms is free software: you can redistribute it and/or modify
7-
* it under the terms of the GNU General Public License as published by
8-
* the Free Software Foundation, either version 3 of the License, or
9-
* (at your option) any later version.
10-
* <p>
11-
* This program is distributed in the hope that it will be useful,
12-
* but WITHOUT ANY WARRANTY; without even the implied warranty of
13-
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14-
* GNU General Public License for more details.
15-
* <p>
16-
* You should have received a copy of the GNU General Public License
17-
* along with this program. If not, see <http://www.gnu.org/licenses/>.
18-
*/
191
package org.neo4j.graphalgo.similarity;
202

21-
import org.neo4j.graphalgo.core.utils.ExceptionUtil;
22-
import org.neo4j.graphalgo.core.utils.StatementApi;
23-
import org.neo4j.internal.kernel.api.exceptions.EntityNotFoundException;
24-
import org.neo4j.internal.kernel.api.exceptions.InvalidTransactionTypeKernelException;
25-
import org.neo4j.internal.kernel.api.exceptions.KernelException;
26-
import org.neo4j.internal.kernel.api.exceptions.explicitindex.AutoIndexingKernelException;
27-
import org.neo4j.kernel.api.KernelTransaction;
28-
import org.neo4j.kernel.internal.GraphDatabaseAPI;
29-
import org.neo4j.logging.Log;
30-
import org.neo4j.values.storable.Values;
31-
32-
import java.util.ArrayList;
33-
import java.util.Iterator;
34-
import java.util.List;
353
import java.util.stream.Stream;
364

37-
public class SimilarityExporter extends StatementApi {
38-
39-
private final Log log;
40-
private final int propertyId;
41-
private final int relationshipTypeId;
42-
43-
public SimilarityExporter(GraphDatabaseAPI api,
44-
Log log, String relationshipType,
45-
String propertyName) {
46-
super(api);
47-
this.log = log;
48-
propertyId = getOrCreatePropertyId(propertyName);
49-
relationshipTypeId = getOrCreateRelationshipId(relationshipType);
50-
}
51-
52-
public void export(Stream<SimilarityResult> similarityPairs, long batchSize) {
53-
int batches = writeSequential(similarityPairs, batchSize);
54-
log.info("ParallelSimilarityExporter: Batch Size: %d, Batches written - sequentially: %d", batchSize, batches);
55-
}
56-
57-
private void export(SimilarityResult similarityResult) {
58-
applyInTransaction(statement -> {
59-
try {
60-
createRelationship(similarityResult, statement);
61-
} catch (KernelException e) {
62-
ExceptionUtil.throwKernelException(e);
63-
}
64-
return null;
65-
});
66-
67-
}
68-
69-
private void export(List<SimilarityResult> similarityResults) {
70-
applyInTransaction(statement -> {
71-
for (SimilarityResult similarityResult : similarityResults) {
72-
try {
73-
createRelationship(similarityResult, statement);
74-
} catch (KernelException e) {
75-
ExceptionUtil.throwKernelException(e);
76-
}
77-
}
78-
return null;
79-
});
80-
81-
}
82-
83-
private void createRelationship(SimilarityResult similarityResult, KernelTransaction statement) throws EntityNotFoundException, InvalidTransactionTypeKernelException, AutoIndexingKernelException {
84-
long node1 = similarityResult.item1;
85-
long node2 = similarityResult.item2;
86-
long relationshipId = statement.dataWrite().relationshipCreate(node1, relationshipTypeId, node2);
87-
88-
statement.dataWrite().relationshipSetProperty(
89-
relationshipId, propertyId, Values.doubleValue(similarityResult.similarity));
90-
}
91-
92-
private int getOrCreateRelationshipId(String relationshipType) {
93-
return applyInTransaction(stmt -> stmt
94-
.tokenWrite()
95-
.relationshipTypeGetOrCreateForName(relationshipType));
96-
}
97-
98-
private int getOrCreatePropertyId(String propertyName) {
99-
return applyInTransaction(stmt -> stmt
100-
.tokenWrite()
101-
.propertyKeyGetOrCreateForName(propertyName));
102-
}
103-
104-
private int writeSequential(Stream<SimilarityResult> similarityPairs, long batchSize) {
105-
log.info("SimilarityExporter: Writing relationships...");
106-
int[] counter = {0};
107-
if (batchSize == 1) {
108-
similarityPairs.forEach(similarityResult -> {
109-
export(similarityResult);
110-
counter[0]++;
111-
});
112-
} else {
113-
Iterator<SimilarityResult> iterator = similarityPairs.iterator();
114-
do {
115-
List<SimilarityResult> batch = take(iterator, Math.toIntExact(batchSize));
116-
export(batch);
117-
if(batch.size() > 0) {
118-
counter[0]++;
119-
}
120-
} while (iterator.hasNext());
121-
}
122-
123-
return counter[0];
124-
}
125-
126-
127-
private static List<SimilarityResult> take(Iterator<SimilarityResult> iterator, int batchSize) {
128-
List<SimilarityResult> result = new ArrayList<>(batchSize);
129-
while (iterator.hasNext() && batchSize-- > 0) {
130-
result.add(iterator.next());
131-
}
132-
return result;
133-
}
134-
135-
5+
public interface SimilarityExporter {
6+
int export(Stream<SimilarityResult> similarityPairs, long batchSize);
1367
}

algo/src/main/java/org/neo4j/graphalgo/similarity/SimilarityProc.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ Stream<SimilaritySummaryResult> writeAndAggregateResults(Stream<SimilarityResult
161161

162162
} else {
163163
try (ProgressTimer timer = builder.timeWrite()) {
164-
SimilarityExporter similarityExporter = new SimilarityExporter(api, log, writeRelationshipType, writeProperty);
164+
SequentialSimilarityExporter similarityExporter = new SequentialSimilarityExporter(api, log, writeRelationshipType, writeProperty);
165165
similarityExporter.export(stream.peek(recorder), writeBatchSize);
166166
}
167167
}

0 commit comments

Comments
 (0)