Append filterSql to source variant_to_person in cohort count (#1130)

docs/generated/UNDERLAY_CONFIG.md

@@ -714,7 +714,7 @@ Pointer to SQL that returns entity id - rollup count (= number of related entity
 
 True to skip copying the id-pairs SQL into a new index table, and use the source SQL directly.
 
-Ignored if the [id pairs SQL](#szgroupitemsidpairssqlfile) is undefined.
+When set filter conditions are directly applied on the id-pairs SQL to fetch cohort counts. Ignored if the [id pairs SQL](#szgroupitemsidpairssqlfile) is undefined.
 
 *Default value:* `false`
 

underlay/src/main/java/bio/terra/tanagra/query/bigquery/BQTable.java

@@ -4,6 +4,7 @@
 import com.google.common.collect.ImmutableMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import org.apache.commons.text.StringSubstitutor;
 
 public class BQTable extends SqlTable {
@@ -24,8 +25,14 @@ public BQTable(String sql) {
 
   @Override
   public String render() {
+    return render(null);
+  }
+
+  @Override
+  public String render(String appendToSql) {
+    String appendNotNull = Optional.ofNullable(appendToSql).orElse("");
     if (isRawSql()) {
-      return "(" + sql + ")";
+      return "(" + sql + appendNotNull + ")";
     } else {
       String template = "`${projectId}.${datasetId}`.${tableName}";
       Map<String, String> params =
@@ -34,7 +41,7 @@ public String render() {
               .put("datasetId", datasetId)
               .put("tableName", tableName)
               .build();
-      return StringSubstitutor.replace(template, params);
+      return StringSubstitutor.replace(template, params) + appendNotNull;
     }
   }
 

underlay/src/main/java/bio/terra/tanagra/query/bigquery/translator/filter/BQHierarchyHasAncestorFilterTranslator.java

@@ -67,6 +67,7 @@ public String buildSql(SqlParams sqlParams, String tableAlias) {
         ancestorDescendantTable.getTablePointer(),
         ancestorIdFilterSql,
         null,
+        false,
         sqlParams,
         hierarchyHasAncestorFilter.getAncestorIds().toArray(new Literal[0]));
   }

underlay/src/main/java/bio/terra/tanagra/query/bigquery/translator/filter/BQHierarchyHasParentFilterTranslator.java

@@ -65,6 +65,7 @@ public String buildSql(SqlParams sqlParams, String tableAlias) {
         childParentIndexTable.getTablePointer(),
         parentIdFilterSql,
         null,
+        false,
         sqlParams);
   }
 

underlay/src/main/java/bio/terra/tanagra/query/bigquery/translator/filter/BQRelationshipFilterTranslator.java

@@ -106,6 +106,7 @@ private String foreignKeyOnSelectEntity(SqlParams sqlParams, String tableAlias)
           filterEntityTable.getTablePointer(),
           inSelectFilterSql,
           null,
+          relationshipFilter.getEntityGroup().isUseSourceIdPairsSql(),
           sqlParams);
     }
   }
@@ -182,6 +183,7 @@ private String foreignKeyOnFilterEntity(SqlParams sqlParams, String tableAlias)
             filterEntityTable.getTablePointer(),
             inSelectFilterSql,
             null,
+            relationshipFilter.getEntityGroup().isUseSourceIdPairsSql(),
             sqlParams);
       }
 
@@ -264,6 +266,7 @@ private String intermediateTable(SqlParams sqlParams, String tableAlias) {
         idPairsTable.getEntityIdField(relationshipFilter.getSelectEntity().getName());
     SqlField filterIdIntTable =
         idPairsTable.getEntityIdField(relationshipFilter.getFilterEntity().getName());
+    boolean appendSqlToTable = relationshipFilter.getEntityGroup().isUseSourceIdPairsSql();
 
     if (!relationshipFilter.hasSubFilter()
         && !relationshipFilter.hasGroupByFilter()
@@ -313,6 +316,7 @@ private String intermediateTable(SqlParams sqlParams, String tableAlias) {
           idPairsTable.getTablePointer(),
           subFilterSql.isEmpty() ? null : subFilterSql,
           havingSql.isEmpty() ? null : havingSql,
+          appendSqlToTable,
           sqlParams);
     } else {
       // id IN (SELECT selectId FROM
@@ -342,6 +346,7 @@ private String intermediateTable(SqlParams sqlParams, String tableAlias) {
                 filterEntityTable.getTablePointer(),
                 subFilterSql,
                 null,
+                appendSqlToTable,
                 sqlParams);
       }
 
@@ -353,6 +358,7 @@ private String intermediateTable(SqlParams sqlParams, String tableAlias) {
             idPairsTable.getTablePointer(),
             filterIdInSelectSql.isEmpty() ? null : filterIdInSelectSql,
             null,
+            appendSqlToTable,
             sqlParams);
       }
 
@@ -374,6 +380,7 @@ private String intermediateTable(SqlParams sqlParams, String tableAlias) {
             idPairsTable.getTablePointer(),
             filterIdInSelectSql.isEmpty() ? null : filterIdInSelectSql,
             havingSql,
+            appendSqlToTable,
             sqlParams);
       }
 

underlay/src/main/java/bio/terra/tanagra/query/sql/SqlTable.java

@@ -22,4 +22,6 @@ public String getSql() {
   }
 
   public abstract String render();
+
+  public abstract String render(String append);
 }

underlay/src/main/java/bio/terra/tanagra/query/sql/translator/ApiTranslator.java

@@ -205,16 +205,18 @@ default String inSelectFilterSql(
       SqlTable table,
       @Nullable String filterSql,
       @Nullable String havingSql,
+      boolean appendSqlToTable,
       SqlParams sqlParams,
       Literal... unionAllLiterals) {
     List<String> selectSqls = new ArrayList<>();
+    String sqlClause =
+        (filterSql != null ? " WHERE " + filterSql : "")
+            + (havingSql != null ? ' ' + havingSql : "");
     selectSqls.add(
         "SELECT "
             + SqlQueryField.of(selectField).renderForSelect()
             + " FROM "
-            + table.render()
-            + (filterSql != null ? " WHERE " + filterSql : "")
-            + (havingSql != null ? ' ' + havingSql : ""));
+            + (appendSqlToTable ? table.render(sqlClause) : table.render() + sqlClause));
     Arrays.stream(unionAllLiterals)
         .forEach(literal -> selectSqls.add("SELECT @" + sqlParams.addParam("val", literal)));
     return SqlQueryField.of(whereField).renderForWhere(tableAlias)

underlay/src/main/java/bio/terra/tanagra/underlay/Underlay.java

@@ -466,7 +466,12 @@ private static GroupItems fromConfigGroupItems(SZGroupItems szGroupItems, List<E
             szGroupItems.foreignKeyAttributeItemsEntity == null
                 ? null
                 : itemsEntity.getAttribute(szGroupItems.foreignKeyAttributeItemsEntity));
-    return new GroupItems(szGroupItems.name, groupEntity, itemsEntity, groupItemsRelationship);
+    return new GroupItems(
+        szGroupItems.name,
+        szGroupItems.useSourceIdPairsSql,
+        groupEntity,
+        itemsEntity,
+        groupItemsRelationship);
   }
 
   private static CriteriaOccurrence fromConfigCriteriaOccurrence(

underlay/src/main/java/bio/terra/tanagra/underlay/entitymodel/entitygroup/CriteriaOccurrence.java

@@ -37,7 +37,7 @@ public CriteriaOccurrence(
       Relationship primaryCriteriaRelationship,
       Map<String, Set<String>> occurrenceAttributesWithInstanceLevelDisplayHints,
       Map<String, Set<String>> occurrenceAttributesWithRollupInstanceLevelDisplayHints) {
-    super(name);
+    super(name, false);
     this.criteriaEntity = criteriaEntity;
     this.occurrenceEntities = ImmutableList.copyOf(occurrenceEntities);
     this.primaryEntity = primaryEntity;

underlay/src/main/java/bio/terra/tanagra/underlay/entitymodel/entitygroup/EntityGroup.java

@@ -10,15 +10,21 @@ public enum Type {
   }
 
   private final String name;
+  private final boolean useSourceIdPairsSql;
 
-  protected EntityGroup(String name) {
+  protected EntityGroup(String name, boolean useSourceIdPairsSql) {
     this.name = name;
+    this.useSourceIdPairsSql = useSourceIdPairsSql;
   }
 
   public String getName() {
     return name;
   }
 
+  public boolean isUseSourceIdPairsSql() {
+    return useSourceIdPairsSql;
+  }
+
   public abstract Type getType();
 
   public abstract boolean includesEntity(String name);

underlay/src/main/java/bio/terra/tanagra/underlay/entitymodel/entitygroup/GroupItems.java

@@ -10,8 +10,12 @@ public class GroupItems extends EntityGroup {
   private final Relationship groupItemsRelationship;
 
   public GroupItems(
-      String name, Entity groupEntity, Entity itemsEntity, Relationship groupItemsRelationship) {
-    super(name);
+      String name,
+      boolean useSourceIdPairsSql,
+      Entity groupEntity,
+      Entity itemsEntity,
+      Relationship groupItemsRelationship) {
+    super(name, useSourceIdPairsSql);
     this.groupEntity = groupEntity;
     this.itemsEntity = itemsEntity;
     this.groupItemsRelationship = groupItemsRelationship;

underlay/src/main/java/bio/terra/tanagra/underlay/serialization/SZGroupItems.java

@@ -50,6 +50,7 @@ public class SZGroupItems {
       name = "SZGroupItems.useSourceIdPairsSql",
       markdown =
           "True to skip copying the id-pairs SQL into a new index table, and use the source SQL directly.\n\n"
+              + "When set filter conditions are directly applied on the id-pairs SQL to fetch cohort counts. "
               + "Ignored if the [id pairs SQL](${SZGroupItems.idPairsSqlFile}) is undefined.",
       optional = true,
       defaultValue = "false")

underlay/src/main/resources/config/datamapping/aouCT_testonly/entitygroup/variantPerson/entityGroup.json

@@ -4,11 +4,11 @@
   "itemsEntity": "person",
   "idPairsSqlFile":  "idPairs.sql",
   "useSourceIdPairsSql": true,
-  "groupEntityIdFieldName": "variant_id",
+  "groupEntityIdFieldName": "vid",
   "itemsEntityIdFieldName": "flattened_person_id",
   "rollupCountsSql": {
     "sqlFile": "rollupCounts.sql",
-    "entityIdFieldName": "variant_id",
+    "entityIdFieldName": "vid",
     "rollupCountFieldName": "num_persons"
   }
 }

underlay/src/main/resources/config/datamapping/aouCT_testonly/entitygroup/variantPerson/idPairs.sql

@@ -1,3 +1,3 @@
-SELECT DISTINCT vid AS variant_id, flattened_person_id
+SELECT DISTINCT vid, flattened_person_id
 FROM `${omopDataset}.variant_to_person`
 CROSS JOIN UNNEST(person_ids) AS flattened_person_id

underlay/src/main/resources/config/datamapping/aouCT_testonly/entitygroup/variantPerson/rollupCounts.sql

@@ -1,3 +1,3 @@
-SELECT vid AS variant_id, ARRAY_LENGTH(person_ids) AS num_persons
+SELECT vid, ARRAY_LENGTH(person_ids) AS num_persons
 /* Wrap variant_to_person table in a SELECT DISTINCT because there is a duplicate row in the test data. */
 FROM (SELECT DISTINCT vid, person_ids FROM `${omopDataset}.variant_to_person` WHERE REGEXP_CONTAINS(vid, r"{indexIdRegex}"))

underlay/src/test/java/bio/terra/tanagra/query/sql/ApiTranslatorTest.java

@@ -172,7 +172,7 @@ void inSelectFilter() {
         apiTranslator.binaryFilterSql(joinField, BinaryOperator.EQUALS, val, null, sqlParams);
     String sql =
         apiTranslator.inSelectFilterSql(
-            field, tableAlias, joinField, joinTable, joinFilterSql, null, sqlParams);
+            field, tableAlias, joinField, joinTable, joinFilterSql, null, false, sqlParams);
     assertEquals(
         "tableAlias.columnName IN (SELECT joinColumnName FROM `projectId.datasetId`.joinTableName WHERE joinColumnName = @val0)",
         sql);
@@ -186,7 +186,16 @@ void inSelectFilter() {
     Literal val1 = Literal.forInt64(25L);
     sql =
         apiTranslator.inSelectFilterSql(
-            field, tableAlias, joinField, joinTable, joinFilterSql, null, sqlParams, val0, val1);
+            field,
+            tableAlias,
+            joinField,
+            joinTable,
+            joinFilterSql,
+            null,
+            false,
+            sqlParams,
+            val0,
+            val1);
     assertEquals(
         "tableAlias.columnName IN (SELECT joinColumnName FROM `projectId.datasetId`.joinTableName WHERE joinColumnName = @val0 UNION ALL SELECT @val1 UNION ALL SELECT @val2)",
         sql);
@@ -197,7 +206,7 @@ void inSelectFilter() {
     val = Literal.forInt64(38L);
     sql =
         apiTranslator.inSelectFilterSql(
-            field, tableAlias, joinField, joinTable, null, null, sqlParams, val);
+            field, tableAlias, joinField, joinTable, null, null, false, sqlParams, val);
     assertEquals(
         "tableAlias.columnName IN (SELECT joinColumnName FROM `projectId.datasetId`.joinTableName UNION ALL SELECT @val0)",
         sql);
@@ -213,6 +222,7 @@ void inSelectFilter() {
             joinTable,
             null,
             "GROUP BY joinColumnName HAVING COUNT(*) > 1",
+            false,
             sqlParams);
     assertEquals(
         "tableAlias.columnName IN (SELECT joinColumnName FROM `projectId.datasetId`.joinTableName GROUP BY joinColumnName HAVING COUNT(*) > 1)",
@@ -237,7 +247,7 @@ void booleanAndOrFilter() {
         apiTranslator.binaryFilterSql(joinField, BinaryOperator.EQUALS, val2, null, sqlParams);
     String filterSql2 =
         apiTranslator.inSelectFilterSql(
-            field, tableAlias, joinField, joinTable, joinFilterSql, null, sqlParams);
+            field, tableAlias, joinField, joinTable, joinFilterSql, null, false, sqlParams);
 
     String sql =
         apiTranslator.booleanAndOrFilterSql(