extra chunking for object relations

VirtualFlyBrain · Sep 1, 2024 · 73fea97 · 73fea97
1 parent fa4e24d
commit 73fea97
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/src/main/java/ebi/spot/neo4j2owl/N2OProcedure.java b/src/main/java/ebi/spot/neo4j2owl/N2OProcedure.java
@@ -73,10 +73,10 @@ public Stream<N2OReturnValue> exportOWLNodes(@Name("skip") Long skip, @Name("lim
 
 	@SuppressWarnings("unused")
 	@Procedure(mode = Mode.WRITE)
-	public Stream<N2OReturnValue> exportOWLEdges(@Name("relationType") String relationType) {
+	public Stream<N2OReturnValue> exportOWLEdges(@Name("relationType") String relationType, @Name("currentChunk") int currentChunk, @Name("limit") int chunkCount) {
 		logger.resetTimer();
 		N2OExportService exportService = new N2OExportService(db);
-		N2OReturnValue result = exportService.owl2ExportEdges(relationType);
+		N2OReturnValue result = exportService.owl2ExportEdges(relationType, currentChunk, chunkCount);
 		return Stream.of(result);
 	}
 

diff --git a/src/main/java/ebi/spot/neo4j2owl/exporter/N2OExportService.java b/src/main/java/ebi/spot/neo4j2owl/exporter/N2OExportService.java
@@ -120,7 +120,7 @@ public N2OReturnValue owl2ExportNodes(Long skip, Long limit) {
 		return returnValue;
 	}
 
-	public N2OReturnValue owl2ExportEdges(String relationType) {
+	public N2OReturnValue owl2ExportEdges(String relationType, int currentChunk, int chunkCount) {
 		n2OEntityManager = new N2OExportManager();
 		qsls_with_no_matching_properties = new HashSet<>();
 		logger.resetTimer();
@@ -143,7 +143,9 @@ public N2OReturnValue owl2ExportEdges(String relationType) {
 				}
 			}
 			if (relationType == null || relationType.isEmpty() || relationType.equals(OBJECT_PROPERTY)) {
-				for (String rel_qsl : getRelations(OWLObjectProperty.class)) {
+				Set<String> objRelations = getRelations(OWLObjectProperty.class);
+				List<Set<String>> chunks = splitSet(objRelations, chunkCount);
+				for (String rel_qsl : chunks.get(currentChunk)) {
 					addRelation(o, rel_qsl);
 				}
 			}
@@ -492,4 +494,25 @@ private int safeAdd(int number, int increment, int maxValue) {
 		// test < number on integer overflow
 		return (test < number || test > maxValue) ? maxValue : test;
 	}
+
+
+	/**
+	 * Splits given set into given number of chunks.
+	 * @param set to split
+	 * @param numChunks number of chunks
+	 * @return List of subsets
+	 */
+	public static <T> List<Set<T>> splitSet(Set<T> set, int numChunks) {
+		List<Set<T>> chunks = new ArrayList<>();
+		List<T> list = new ArrayList<>(set);
+		int chunkSize = (int) Math.ceil((double) list.size() / numChunks);
+
+		for (int i = 0; i < list.size(); i += chunkSize) {
+			int end = Math.min(list.size(), i + chunkSize);
+			chunks.add(new HashSet<>(list.subList(i, end)));
+		}
+
+		return chunks;
+	}
+
 }