|
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