|
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 | 1 | package org.neo4j.graphalgo.similarity;
|
20 | 2 |
|
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 | 3 | import java.util.stream.Stream;
|
36 | 4 |
|
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); |
136 | 7 | }
|
0 commit comments