|
| 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 | +} |
0 commit comments