Optionally use rollup counts SQL instead of Dataflow job for GroupItems entity group.

annotationProcessor/src/main/java/bio/terra/tanagra/annotation/UnderlayConfigPath.java

@@ -1,20 +1,6 @@
 package bio.terra.tanagra.annotation;
 
-import bio.terra.tanagra.underlay.serialization.SZAttribute;
-import bio.terra.tanagra.underlay.serialization.SZBigQuery;
-import bio.terra.tanagra.underlay.serialization.SZCorePlugin;
-import bio.terra.tanagra.underlay.serialization.SZCriteriaOccurrence;
-import bio.terra.tanagra.underlay.serialization.SZCriteriaSelector;
-import bio.terra.tanagra.underlay.serialization.SZDataType;
-import bio.terra.tanagra.underlay.serialization.SZEntity;
-import bio.terra.tanagra.underlay.serialization.SZGroupItems;
-import bio.terra.tanagra.underlay.serialization.SZHierarchy;
-import bio.terra.tanagra.underlay.serialization.SZIndexer;
-import bio.terra.tanagra.underlay.serialization.SZPrepackagedCriteria;
-import bio.terra.tanagra.underlay.serialization.SZService;
-import bio.terra.tanagra.underlay.serialization.SZTextSearch;
-import bio.terra.tanagra.underlay.serialization.SZUnderlay;
-import bio.terra.tanagra.underlay.serialization.SZVisualization;
+import bio.terra.tanagra.underlay.serialization.*;
 import java.util.List;
 
 @SuppressWarnings("PMD.CouplingBetweenObjects")
@@ -52,6 +38,7 @@ public class UnderlayConfigPath extends AnnotationPath {
           SZCriteriaSelector.Display.class,
           SZCriteriaSelector.Modifier.class,
           SZPrepackagedCriteria.class,
+          SZRollupCountsSql.class,
           SZVisualization.class,
           SZCorePlugin.class);
 

docs/generated/UNDERLAY_CONFIG.md

@@ -23,6 +23,7 @@ This documentation is generated from annotations in the configuration classes.
 * [SZPrepackagedCriteria](#szprepackagedcriteria)
 * [SZPrimaryCriteriaRelationship](#szprimarycriteriarelationship)
 * [SZPrimaryRelationship](#szprimaryrelationship)
+* [SZRollupCountsSql](#szrollupcountssql)
 * [SZService](#szservice)
 * [SZSourceData](#szsourcedata)
 * [SZSourceQuery](#szsourcequery)
@@ -642,7 +643,7 @@ Required if the [id pairs SQL](#szgroupitemsidpairssqlfile) is defined.
 
 Name of the group entity - items entity id pairs SQL file.
 
-If this property is set, then the [id pairs SQL](#szgroupitemsidpairssqlfile) must be unset. File must be in the same directory as the entity group file. Name includes file extension.
+File must be in the same directory as the entity group file. Name includes file extension.
 
 There can be other columns selected in the SQL file (e.g. `SELECT * FROM relationships`), but the group and items entity ids are required. If this property is set, then the [foreign key atttribute](#szgroupitemsforeignkeyattributeitemsentity) must be unset.
 
@@ -669,6 +670,11 @@ Name of the entity group.
 
 This is the unique identifier for the entity group. In a single underlay, the entity group names of any group type cannot overlap. Name may not include spaces or special characters, only letters and numbers. The first character must be a letter.
 
+### SZGroupItems.rollupCountsSql
+**optional** [SZRollupCountsSql](#szrollupcountssql)
+
+Pointer to SQL that returns entity id - rollup count (= number of related entity instances) pairs.
+
 ### SZGroupItems.useSourceIdPairsSql
 **optional** boolean
 
@@ -967,6 +973,36 @@ Name of the field or column name that maps to the primary entity id. Required if
 
 
 
+## SZRollupCountsSql
+Pointer to SQL that returns entity id - rollup count (= number of related entity instances) pairs (e.g. variant - number of people). Useful when there's an easy way to calculate these in SQL and we want to avoid ingesting the full entity - related entity relationship id pairs table into Dataflow.
+
+### SZRollupCountsSql.entityIdFieldName
+**required** String
+
+Name of the field or column name that maps to the entity id.
+
+*Example value:* `entity_id`
+
+### SZRollupCountsSql.rollupCountFieldName
+**required** String
+
+Name of the field or column name that maps to the rollup count per entity id.
+
+*Example value:* `rollup_count`
+
+### SZRollupCountsSql.rollupCountsSqlFile
+**required** String
+
+Name of the entity id - rollup counts (= number of items entity instances) pairs SQL file.
+
+File must be in the same directory as the entity/group file. Name includes file extension.
+
+There can be other columns selected in the SQL file (e.g. `SELECT * FROM relationships`), but the entity id and rollup count fields are required.
+
+*Example value:* `rollupCounts.sql`
+
+
+
 ## SZService
 Service configuration.
 

indexer/src/main/java/bio/terra/tanagra/indexing/JobSequencer.java

@@ -31,11 +31,7 @@
 import bio.terra.tanagra.underlay.indextable.ITHierarchyChildParent;
 import bio.terra.tanagra.underlay.indextable.ITRelationshipIdPairs;
 import bio.terra.tanagra.underlay.serialization.SZIndexer;
-import bio.terra.tanagra.underlay.sourcetable.STEntityAttributes;
-import bio.terra.tanagra.underlay.sourcetable.STHierarchyChildParent;
-import bio.terra.tanagra.underlay.sourcetable.STHierarchyRootFilter;
-import bio.terra.tanagra.underlay.sourcetable.STRelationshipIdPairs;
-import bio.terra.tanagra.underlay.sourcetable.STTextSearchTerms;
+import bio.terra.tanagra.underlay.sourcetable.*;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -192,6 +188,14 @@ public static SequencedJobSet getJobSetForGroupItems(
                     groupItems.getGroupEntity().getName(),
                     groupItems.getItemsEntity().getName())
             : null;
+    STRelationshipRollupCounts groupItemsRelationshipRollupCountsTable =
+        underlay
+            .getSourceSchema()
+            .getRelationshipRollupCounts(
+                groupItems.getName(),
+                groupItems.getGroupEntity().getName(),
+                groupItems.getItemsEntity().getName())
+            .orElse(null);
     jobSet.addJob(
         new WriteRollupCounts(
             indexerConfig,
@@ -203,7 +207,8 @@ public static SequencedJobSet getJobSetForGroupItems(
             itemsEntityIndexTable,
             groupItemsIdPairsTable,
             null,
-            null));
+            null,
+            groupItemsRelationshipRollupCountsTable));
 
     // If the criteria entity has hierarchies, then also compute the criteria rollup counts for each
     // hierarchy.
@@ -228,7 +233,8 @@ public static SequencedJobSet getJobSetForGroupItems(
                           underlay
                               .getIndexSchema()
                               .getHierarchyAncestorDescendant(
-                                  groupItems.getGroupEntity().getName(), hierarchy.getName()))));
+                                  groupItems.getGroupEntity().getName(), hierarchy.getName()),
+                          null)));
     }
 
     if (groupItems.getGroupEntity().hasHierarchies()) {
@@ -378,6 +384,7 @@ public static SequencedJobSet getJobSetForCriteriaOccurrence(
             primaryEntityIndexTable,
             primaryCriteriaIdPairsTable,
             null,
+            null,
             null));
 
     // If the criteria entity has hierarchies, then also compute the criteria rollup counts for each
@@ -404,7 +411,8 @@ public static SequencedJobSet getJobSetForCriteriaOccurrence(
                               .getIndexSchema()
                               .getHierarchyAncestorDescendant(
                                   criteriaOccurrence.getCriteriaEntity().getName(),
-                                  hierarchy.getName()))));
+                                  hierarchy.getName()),
+                          null)));
     }
 
     // Compute instance-level display hints for the occurrence entity attributes that are flagged as

indexer/src/main/java/bio/terra/tanagra/indexing/job/dataflow/WriteRollupCounts.java

@@ -17,6 +17,8 @@
 import bio.terra.tanagra.underlay.indextable.ITHierarchyAncestorDescendant;
 import bio.terra.tanagra.underlay.indextable.ITRelationshipIdPairs;
 import bio.terra.tanagra.underlay.serialization.SZIndexer;
+import bio.terra.tanagra.underlay.sourcetable.*;
+import com.google.api.*;
 import com.google.api.services.bigquery.model.TableFieldSchema;
 import com.google.api.services.bigquery.model.TableRow;
 import com.google.api.services.bigquery.model.TableSchema;
@@ -78,6 +80,7 @@ public class WriteRollupCounts extends BigQueryJob {
   private final @Nullable ITRelationshipIdPairs relationshipIdPairsIndexTable;
   private final @Nullable Hierarchy hierarchy;
   private final @Nullable ITHierarchyAncestorDescendant ancestorDescendantTable;
+  private final @Nullable STRelationshipRollupCounts relationshipRollupCountsSourceTable;
 
   @SuppressWarnings({"checkstyle:ParameterNumber", "PMD.ExcessiveParameterList"})
   public WriteRollupCounts(
@@ -90,7 +93,8 @@ public WriteRollupCounts(
       ITEntityMain countedEntityIndexTable,
       @Nullable ITRelationshipIdPairs relationshipIdPairsIndexTable,
       @Nullable Hierarchy hierarchy,
-      @Nullable ITHierarchyAncestorDescendant ancestorDescendantTable) {
+      @Nullable ITHierarchyAncestorDescendant ancestorDescendantTable,
+      @Nullable STRelationshipRollupCounts relationshipRollupCountsSourceTable) {
     super(indexerConfig);
     this.entityGroup = entityGroup;
     this.entity = entity;
@@ -101,6 +105,7 @@ public WriteRollupCounts(
     this.relationshipIdPairsIndexTable = relationshipIdPairsIndexTable;
     this.hierarchy = hierarchy;
     this.ancestorDescendantTable = ancestorDescendantTable;
+    this.relationshipRollupCountsSourceTable = relationshipRollupCountsSourceTable;
   }
 
   @Override
@@ -136,16 +141,20 @@ && outputTableHasAtLeastOneRowWithNotNullField(
 
   @Override
   public void run(boolean isDryRun) {
-    // Only run the Dataflow job if the temp table hasn't been written yet.
-    Optional<Table> tempTable =
-        googleBigQuery.getTable(
-            indexerConfig.bigQuery.indexData.projectId,
-            indexerConfig.bigQuery.indexData.datasetId,
-            getTempTableName());
-    if (tempTable.isEmpty()) {
-      writeFieldsToTempTable(isDryRun);
+    if (relationshipRollupCountsSourceTable == null) {
+      // Only run the Dataflow job if the temp table hasn't been written yet.
+      Optional<Table> tempTable =
+          googleBigQuery.getTable(
+              indexerConfig.bigQuery.indexData.projectId,
+              indexerConfig.bigQuery.indexData.datasetId,
+              getTempTableName());
+      if (tempTable.isEmpty()) {
+        writeFieldsToTempTable(isDryRun);
+      } else {
+        LOGGER.info("Temp table has already been written. Skipping Dataflow job.");
+      }
     } else {
-      LOGGER.info("Temp table has already been written. Skipping Dataflow job.");
+      LOGGER.info("Rollup counts source SQL is defined. Skipping Dataflow job.");
     }
 
     // Dataflow jobs can only write new rows to BigQuery, so in this second step, copy over the
@@ -353,21 +362,30 @@ private void copyFieldsToEntityTable(boolean isDryRun) {
         indexTable.getEntityGroupCountField(
             entityGroup.getName(), hierarchy == null ? null : hierarchy.getName());
 
-    BQTable tempBQTable =
-        new BQTable(
-            indexerConfig.bigQuery.indexData.projectId,
-            indexerConfig.bigQuery.indexData.datasetId,
-            getTempTableName());
-    SqlField tempTableIdField = SqlField.of(entityTableIdField.getColumnName());
-    SqlField tempTableCountField = SqlField.of(entityTableCountField.getColumnName());
-    String tempTableSql =
+    BQTable rollupCountsBQTable;
+    SqlField rollupCountsTableIdField;
+    SqlField rollupCountsTableCountField;
+    if (relationshipRollupCountsSourceTable == null) {
+      rollupCountsBQTable =
+          new BQTable(
+              indexerConfig.bigQuery.indexData.projectId,
+              indexerConfig.bigQuery.indexData.datasetId,
+              getTempTableName());
+      rollupCountsTableIdField = SqlField.of(entityTableIdField.getColumnName());
+      rollupCountsTableCountField = SqlField.of(entityTableCountField.getColumnName());
+    } else {
+      rollupCountsBQTable = relationshipRollupCountsSourceTable.getTablePointer();
+      rollupCountsTableIdField = relationshipRollupCountsSourceTable.getEntityIdField();
+      rollupCountsTableCountField = relationshipRollupCountsSourceTable.getCountField();
+    }
+    String rollupCountsTableSql =
         "SELECT "
-            + SqlQueryField.of(tempTableIdField).renderForSelect()
+            + SqlQueryField.of(rollupCountsTableIdField).renderForSelect()
             + ", "
-            + SqlQueryField.of(tempTableCountField).renderForSelect()
+            + SqlQueryField.of(rollupCountsTableCountField).renderForSelect()
             + " FROM "
-            + tempBQTable.render();
-    LOGGER.info("temp table query: {}", tempTableSql);
+            + rollupCountsBQTable.render();
+    LOGGER.info("rollup counts table query: {}", rollupCountsTableSql);
 
     // Build an update-from-select query for the index entity main table and the
     // id-count query.
@@ -381,19 +399,19 @@ private void copyFieldsToEntityTable(boolean isDryRun) {
             + " SET "
             + SqlQueryField.of(entityTableCountField).renderForSelect(updateTableAlias)
             + " = "
-            + SqlQueryField.of(tempTableCountField).renderForSelect(tempTableAlias)
+            + SqlQueryField.of(rollupCountsTableCountField).renderForSelect(tempTableAlias)
             + " FROM (SELECT "
-            + SqlQueryField.of(tempTableCountField).renderForSelect()
+            + SqlQueryField.of(rollupCountsTableCountField).renderForSelect()
             + ", "
-            + SqlQueryField.of(tempTableIdField).renderForSelect()
+            + SqlQueryField.of(rollupCountsTableIdField).renderForSelect()
             + " FROM "
-            + tempBQTable.render()
+            + rollupCountsBQTable.render()
             + ") AS "
             + tempTableAlias
             + " WHERE "
             + SqlQueryField.of(entityTableIdField).renderForSelect(updateTableAlias)
             + " = "
-            + SqlQueryField.of(tempTableIdField).renderForSelect(tempTableAlias);
+            + SqlQueryField.of(rollupCountsTableIdField).renderForSelect(tempTableAlias);
     LOGGER.info("update-from-select query: {}", updateFromSelectSql);
 
     // Run the update-from-select to write the count field in the index entity main table.

ui/src/tanagra-underlay/underlayConfig.ts

@@ -114,6 +114,7 @@ export type SZGroupItems = {
   itemsEntity: string;
   itemsEntityIdFieldName?: string;
   name: string;
+  rollupCountsSql?: SZRollupCountsSql;
   useSourceIdPairsSql?: boolean;
 };
 
@@ -176,6 +177,12 @@ export type SZPrimaryRelationship = {
   primaryEntityIdFieldName?: string;
 };
 
+export type SZRollupCountsSql = {
+  entityIdFieldName: string;
+  rollupCountFieldName: string;
+  sqlFile: string;
+};
+
 export type SZService = {
   bigQuery: SZBigQuery;
   underlay: string;