Parallelize per tree element

ingef · Feb 22, 2024 · b23898d · b23898d
1 parent 650243e
commit b23898d
Show file tree

Hide file tree

Showing 7 changed files with 68 additions and 24 deletions.
diff --git a/backend/src/main/java/com/bakdata/conquery/models/config/SqlConnectorConfig.java b/backend/src/main/java/com/bakdata/conquery/models/config/SqlConnectorConfig.java
@@ -1,5 +1,7 @@
 package com.bakdata.conquery.models.config;
 
+import javax.validation.constraints.Min;
+
 import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.Data;
@@ -29,4 +31,10 @@ public class SqlConnectorConfig {
 	private String jdbcConnectionUrl;
 
 	private String primaryColumn = DEFAULT_PRIMARY_COLUMN;
+
+	/**
+	 * The amount of threads for background tasks like calculating matching stats {@link com.bakdata.conquery.models.jobs.SqlUpdateMatchingStatsJob}.
+	 */
+	@Min(1)
+	private int backgroundThreads;
 }
diff --git a/...end/src/main/java/com/bakdata/conquery/models/datasets/concepts/tree/ConceptTreeNode.java b/...end/src/main/java/com/bakdata/conquery/models/datasets/concepts/tree/ConceptTreeNode.java
@@ -27,5 +27,7 @@ public interface ConceptTreeNode<ID extends ConceptElementId<? extends ConceptEl
 
 	MatchingStats getMatchingStats();
 
+	void setMatchingStats(MatchingStats matchingStats);
+
 	String getDescription();
 }
diff --git a/backend/src/main/java/com/bakdata/conquery/models/jobs/SqlUpdateMatchingStatsJob.java b/backend/src/main/java/com/bakdata/conquery/models/jobs/SqlUpdateMatchingStatsJob.java
@@ -8,9 +8,13 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.function.BinaryOperator;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import com.bakdata.conquery.mode.local.SqlMatchingStats;
 import com.bakdata.conquery.models.common.daterange.CDateRange;
@@ -26,6 +30,7 @@
 import com.bakdata.conquery.sql.conversion.model.ColumnDateRange;
 import com.bakdata.conquery.sql.execution.SqlExecutionService;
 import com.bakdata.conquery.util.TablePrimaryColumnUtil;
+import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.jooq.Condition;
 import org.jooq.DSLContext;
@@ -55,6 +60,7 @@ public class SqlUpdateMatchingStatsJob extends Job {
 	private final DSLContext dslContext;
 	private final SqlFunctionProvider functionProvider;
 	private final Collection<Concept<?>> concepts;
+	private final ExecutorService executors;
 
 	public SqlUpdateMatchingStatsJob(
 			SqlConnectorConfig sqlConnectorConfig,
@@ -67,6 +73,7 @@ public SqlUpdateMatchingStatsJob(
 		this.dslContext = executionService.getDslContext();
 		this.functionProvider = functionProvider;
 		this.concepts = concepts;
+		this.executors = Executors.newFixedThreadPool(sqlConnectorConfig.getBackgroundThreads());
 	}
 
 	@Override
@@ -79,34 +86,38 @@ public void execute() throws Exception {
 
 		log.debug("BEGIN update Matching stats for {} Concepts.", concepts.size());
 
-		for (Concept<?> concept : concepts) {
-			if (!(concept instanceof TreeConcept tree)) {
-				log.error("Collecting MatchingStats is currently only supported for TreeConcepts.");
-				break;
-			}
-			SqlMatchingStats matchingStats = collectMatchingStats(concept.getConnectors(), tree);
-			concept.setMatchingStats(matchingStats);
+		concepts.stream()
+				.filter(SqlUpdateMatchingStatsJob::isTreeConcept)
+				.flatMap(concept -> collectMatchingStats(concept.getConnectors(), (TreeConcept) concept))
+				.forEach(executors::submit);
+
+		executors.shutdown();
+		while (!executors.awaitTermination(1, TimeUnit.MINUTES)) {
+			log.debug("Waiting for executors to set matching stats for all concepts...");
 		}
 
 		log.debug("DONE collecting matching stats.");
 	}
 
-	private SqlMatchingStats collectMatchingStats(List<? extends Connector> connectors, ConceptTreeNode<?> treeNode) {
-
-		treeNode.getChildren().forEach(child -> {
-			SqlMatchingStats childStats = collectMatchingStats(connectors, child);
-			child.setMatchingStats(childStats);
-		});
-
-		Optional<CTCondition> childCondition = treeNode instanceof ConceptTreeChild treeChild
-											   ? Optional.of(treeChild.getCondition())
-											   : Optional.empty();
+	@Override
+	public void cancel() {
+		super.cancel();
+		executors.shutdownNow();
+	}
 
-		long events = collectEventCount(connectors, childCondition);
-		long entities = collectEntityCount(connectors, childCondition);
-		CDateRange span = collectDateSpan(connectors, childCondition);
+	private static boolean isTreeConcept(Concept<?> concept) {
+		if (!(concept instanceof TreeConcept)) {
+			log.error("Collecting MatchingStats is currently only supported for TreeConcepts.");
+			return false;
+		}
+		return true;
+	}
 
-		return new SqlMatchingStats(events, entities, span);
+	private Stream<Runnable> collectMatchingStats(List<? extends Connector> connectors, ConceptTreeNode<?> treeNode) {
+		return Stream.concat(
+				treeNode.getChildren().stream().flatMap(child -> collectMatchingStats(connectors, child)),
+				Stream.of(new SqlMatchingStatsTask(connectors, treeNode))
+		);
 	}
 
 	/**
@@ -261,4 +272,25 @@ private Condition toJooqCondition(Connector connector, Optional<CTCondition> chi
 							 .orElse(DSL.noCondition());
 	}
 
+	@RequiredArgsConstructor
+	private class SqlMatchingStatsTask implements Runnable {
+
+		private final List<? extends Connector> connectors;
+		private final ConceptTreeNode<?> treeNode;
+
+		@Override
+		public void run() {
+			Optional<CTCondition> childCondition = treeNode instanceof ConceptTreeChild treeChild
+												   ? Optional.of(treeChild.getCondition())
+												   : Optional.empty();
+
+			long events = collectEventCount(connectors, childCondition);
+			long entities = collectEntityCount(connectors, childCondition);
+			CDateRange span = collectDateSpan(connectors, childCondition);
+
+			SqlMatchingStats matchingStats = new SqlMatchingStats(events, entities, span);
+			treeNode.setMatchingStats(matchingStats);
+		}
+	}
+
 }
diff --git a/...nd/src/main/java/com/bakdata/conquery/sql/conversion/dialect/HanaSqlFunctionProvider.java b/...nd/src/main/java/com/bakdata/conquery/sql/conversion/dialect/HanaSqlFunctionProvider.java
@@ -78,7 +78,7 @@ public ColumnDateRange daterange(CDateRange dateRestriction) {
 	}
 
 	@Override
-	public ColumnDateRange daterange(ValidityDate validityDate, String qualifier, String conceptLabel) {
+	public ColumnDateRange daterange(ValidityDate validityDate, String qualifier, String label) {
 
 		Column startColumn;
 		Column endColumn;
@@ -104,7 +104,7 @@ public ColumnDateRange daterange(ValidityDate validityDate, String qualifier, St
 		);
 
 		return ColumnDateRange.of(rangeStart, rangeEnd)
-							  .asValidityDateRange(conceptLabel);
+							  .asValidityDateRange(label);
 	}
 
 	@Override

diff --git a/...d/src/test/java/com/bakdata/conquery/integration/sql/dialect/HanaSqlIntegrationTests.java b/...d/src/test/java/com/bakdata/conquery/integration/sql/dialect/HanaSqlIntegrationTests.java
@@ -176,6 +176,7 @@ public HanaTestcontainerContextProvider() {
 														.databaseUsername(hanaContainer.getUsername())
 														.databasePassword(hanaContainer.getPassword())
 														.primaryColumn("pid")
+														.backgroundThreads(Runtime.getRuntime().availableProcessors())
 														.build();
 			this.dslContext = DslContextFactory.create(sqlConnectorConfig);
 		}

diff --git a/...rc/test/java/com/bakdata/conquery/integration/sql/dialect/PostgreSqlIntegrationTests.java b/...rc/test/java/com/bakdata/conquery/integration/sql/dialect/PostgreSqlIntegrationTests.java
@@ -70,6 +70,7 @@ static void before() {
 									  .databasePassword(PASSWORD)
 									  .withPrettyPrinting(true)
 									  .primaryColumn("pid")
+									  .backgroundThreads(Runtime.getRuntime().availableProcessors())
 									  .build();
 		dslContext = DslContextFactory.create(sqlConfig);
 		testSqlDialect = new TestPostgreSqlDialect(dslContext);

diff --git a/backend/src/test/resources/tests/matchingstats/kh-content.csv b/backend/src/test/resources/tests/matchingstats/kh-content.csv
@@ -1,4 +1,4 @@
-pid,icd_code,aufnahmedatum,entlassungsdatum
+primary_column,icd_code,aufnahmedatum,entlassungsdatum
 3,"F200",2022-11-28,2022-11-11
 3,"F201",2021-08-31,2021-12-15
 3,"F204",2010-07-01,2019-07-13