Pertaining cohort covariates

pom.xml

@@ -142,7 +142,7 @@
     <dependency>
       <groupId>org.ohdsi.sql</groupId>
       <artifactId>SqlRender</artifactId>
-      <version>1.6.8</version>
+      <version>1.16.1-SNAPSHOT</version>
       <scope>test</scope>
     </dependency>
     <dependency>

src/main/java/org/ohdsi/circe/cohortdefinition/ConditionEra.java

@@ -19,7 +19,16 @@
 package org.ohdsi.circe.cohortdefinition;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+
+import org.apache.commons.lang3.StringUtils;
 import org.ohdsi.circe.cohortdefinition.builders.BuilderOptions;
+import org.ohdsi.circe.cohortdefinition.builders.ColumnFieldData;
+import org.ohdsi.circe.cohortdefinition.builders.ColumnFieldDataType;
 import org.ohdsi.circe.vocabulary.Concept;
 
 /**
@@ -61,4 +70,62 @@ public String accept(IGetCriteriaSqlDispatcher dispatcher, BuilderOptions option
     return dispatcher.getCriteriaSql(this, options);
   }
 
+  @Override
+  public List<ColumnFieldData> getSelectedField(Boolean retainCohortCovariates) {
+      List<ColumnFieldData> selectCols = new ArrayList<>();
+
+      if (retainCohortCovariates) {
+          if (occurrenceCount != null) {
+              selectCols.add(new ColumnFieldData("condition_occurrence_count", ColumnFieldDataType.INTEGER));
+          }
+      }
+
+      return selectCols;
+  }
+
+  @Override
+  public String embedWindowedCriteriaQuery(String query, Map<String, ColumnFieldData> mapDistinctField) {
+      List<String> selectCols = new ArrayList<>();
+      List<String> groupCols = new ArrayList<>();
+      for (Entry<String, ColumnFieldData> entry : mapDistinctField.entrySet()) {
+          if (entry.getKey().equals("condition_occurrence_count") && occurrenceCount != null) {
+              selectCols.add(", cc.condition_occurrence_count");
+              groupCols.add(", cc.condition_occurrence_count");
+          } else {
+              selectCols.add(", CAST(null as " + entry.getValue().getDataType().getType() + ") " + entry.getKey());
+          }
+      }
+
+      query = StringUtils.replace(query, "@additionColumnscc", StringUtils.join(selectCols, ""));
+      query = StringUtils.replace(query, "@additionGroupColumnscc", StringUtils.join(groupCols, ""));
+      return query;
+  }
+
+
+  @Override
+  public String embedWindowedCriteriaQueryP(String query) {
+      ArrayList<String> selectColsA = new ArrayList<>();
+
+      if (occurrenceCount != null) {
+          selectColsA.add(", A.condition_occurrence_count");
+      }
+
+      query = StringUtils.replace(query, "@p.additionColumns", StringUtils.join(selectColsA, ""));
+      return query;
+  }
+
+  @Override
+  public String embedWrapCriteriaQuery(String query, List<String> selectColsPE, BuilderOptions options) {
+      ArrayList<String> selectCols = new ArrayList<>();
+
+      if(!options.isPrimaryCriteria()){
+        if (occurrenceCount != null) {
+            selectCols.add(", Q.condition_occurrence_count");
+            selectColsPE.add(", PE.condition_occurrence_count");
+        }
+      }
+
+      query = StringUtils.replace(query, "@QAdditionalColumnsInclusionN", StringUtils.join(selectCols, ""));
+      return query;
+  }
 }

src/main/java/org/ohdsi/circe/cohortdefinition/ConditionOccurrence.java

@@ -20,7 +20,16 @@
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonTypeName;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+
+import org.apache.commons.lang3.StringUtils;
 import org.ohdsi.circe.cohortdefinition.builders.BuilderOptions;
+import org.ohdsi.circe.cohortdefinition.builders.ColumnFieldData;
+import org.ohdsi.circe.cohortdefinition.builders.ColumnFieldDataType;
 import org.ohdsi.circe.vocabulary.Concept;
 
 /**
@@ -75,4 +84,113 @@ public String accept(IGetCriteriaSqlDispatcher dispatcher, BuilderOptions option
   {
     return dispatcher.getCriteriaSql(this, options);
   }
+
+  @Override
+  public List<ColumnFieldData> getSelectedField(Boolean retainCohortCovariates) {
+      List<ColumnFieldData> selectCols = new ArrayList<>();
+
+      if (retainCohortCovariates) {
+          if (conditionType != null && conditionType.length > 0) {
+              selectCols.add(new ColumnFieldData("condition_type_concept_id", ColumnFieldDataType.INTEGER));
+          }
+
+          if (conditionSourceConcept != null) {
+              selectCols.add(new ColumnFieldData("condition_source_concept_id", ColumnFieldDataType.INTEGER));
+          }
+
+          // providerSpecialty
+          if (providerSpecialty != null && providerSpecialty.length > 0) {
+              selectCols.add(new ColumnFieldData("provider_id", ColumnFieldDataType.INTEGER));
+          }
+
+          if (conditionStatus != null && conditionStatus.length > 0) {
+              selectCols.add(new ColumnFieldData("condition_status_concept_id", ColumnFieldDataType.INTEGER));
+          }
+      }
+
+      return selectCols;
+  }
+
+  @Override
+  public String embedWindowedCriteriaQuery(String query, Map<String, ColumnFieldData> mapDistinctField) {
+      List<String> selectCols = new ArrayList<>();
+      List<String> groupCols = new ArrayList<>();
+      for (Entry<String, ColumnFieldData> entry : mapDistinctField.entrySet()) {
+          if (entry.getKey().equals("condition_type_concept_id") && conditionType != null && conditionType.length > 0) {
+              selectCols.add(", cc.condition_type_concept_id");
+              groupCols.add(", cc.condition_type_concept_id");
+          } else if (entry.getKey().equals("condition_source_concept_id") && conditionSourceConcept != null) {
+              selectCols.add(", cc.condition_source_concept_id");
+              groupCols.add(", cc.condition_source_concept_id");
+          } else if (entry.getKey().equals("provider_id") && providerSpecialty != null && providerSpecialty.length > 0) {
+              selectCols.add(", cc.provider_id");
+              groupCols.add(", cc.provider_id");
+          } else if (entry.getKey().equals("condition_status_concept_id") && conditionStatus != null && conditionStatus.length > 0) {
+              selectCols.add(", cc.condition_status_concept_id");
+              groupCols.add(", cc.condition_status_concept_id");
+          } else {
+              selectCols.add(", CAST(null as " + entry.getValue().getDataType().getType() + ") " + entry.getKey());
+          }
+      }
+
+      query = StringUtils.replace(query, "@additionColumnscc", StringUtils.join(selectCols, ""));
+      query = StringUtils.replace(query, "@additionGroupColumnscc", StringUtils.join(groupCols, ""));
+      return query;
+  }
+
+  @Override
+  public String embedWindowedCriteriaQueryP(String query) {
+      ArrayList<String> selectColsA = new ArrayList<>();
+
+      if (conditionType != null && conditionType.length > 0) {
+          selectColsA.add(", A.condition_type_concept_id");
+      }
+
+      if (conditionSourceConcept != null) {
+          selectColsA.add(", A.condition_source_concept_id");
+      }
+
+      // providerSpecialty
+      if (providerSpecialty != null && providerSpecialty.length > 0) {
+          selectColsA.add(", A.provider_id");
+      }
+
+      if (conditionStatus != null && conditionStatus.length > 0) {
+          selectColsA.add(", A.condition_status_concept_id");
+      }
+
+      query = StringUtils.replace(query, "@p.additionColumns", StringUtils.join(selectColsA, ""));
+      return query;
+  }
+
+  @Override
+  public String embedWrapCriteriaQuery(String query, List<String> selectColsPE, BuilderOptions options) {
+      ArrayList<String> selectCols = new ArrayList<>();
+
+      if(!options.isPrimaryCriteria()){
+        if (conditionType != null && conditionType.length > 0) {
+            selectCols.add(", Q.condition_type_concept_id");
+            selectColsPE.add(", PE.condition_type_concept_id");
+        }
+
+        if (conditionSourceConcept != null) {
+            selectCols.add(", Q.condition_source_concept_id");
+            selectColsPE.add(", PE.condition_source_concept_id");
+        }
+
+        // providerSpecialty
+        if (providerSpecialty != null && providerSpecialty.length > 0) {
+            selectCols.add(", Q.provider_id");
+            selectColsPE.add(", PE.provider_id");
+        }
+
+        if (conditionStatus != null && conditionStatus.length > 0) {
+            selectCols.add(", Q.condition_status_concept_id");
+            selectColsPE.add(", PE.condition_status_concept_id");
+        }
+      }
+
+      query = StringUtils.replace(query, "@QAdditionalColumnsInclusionN", StringUtils.join(selectCols, ""));
+      return query;
+  }
 }

src/main/java/org/ohdsi/circe/cohortdefinition/Criteria.java

@@ -22,7 +22,16 @@
 import com.fasterxml.jackson.annotation.JsonSubTypes;
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import com.fasterxml.jackson.annotation.JsonTypeInfo.As;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.commons.lang3.StringUtils;
 import org.ohdsi.circe.cohortdefinition.builders.BuilderOptions;
+import org.ohdsi.circe.cohortdefinition.builders.ColumnFieldData;
+import org.ohdsi.circe.cohortdefinition.builders.ColumnFieldDataType;
 
 /**
  *
@@ -54,6 +63,25 @@ public String accept(IGetCriteriaSqlDispatcher dispatcher) {
   }
 
   public abstract String accept(IGetCriteriaSqlDispatcher dispatcher, BuilderOptions options);
+
+  public List<ColumnFieldData> getSelectedField(Boolean retainCohortCovariates) {
+      return new ArrayList<>();
+  }
+
+  public String embedWindowedCriteriaQuery(String query, Map<String, ColumnFieldData> mapDistinctField) {
+      query = StringUtils.replace(query, "@additionColumnscc", "");
+      return query;
+  }
+
+  public String embedWindowedCriteriaQueryP(String query) {
+      query = StringUtils.replace(query, "@p.additionColumns", "");
+      return query;
+  }
+
+  public String embedWrapCriteriaQuery(String query, List<String> selectColsPE, BuilderOptions options) {
+      query = StringUtils.replace(query, "@QAdditionalColumnsInclusionN", "");
+      return query;
+  }
 
   @JsonProperty("CorrelatedCriteria")
   public CriteriaGroup CorrelatedCriteria;

src/main/java/org/ohdsi/circe/cohortdefinition/CustomEraStrategy.java

@@ -39,7 +39,7 @@ public class CustomEraStrategy  extends EndStrategy {
   public Integer daysSupplyOverride = null;
 
   @Override
-  public String accept(IGetEndStrategySqlDispatcher dispatcher, String eventTable) {
-    return dispatcher.getStrategySql(this, eventTable);
+  public String accept(IGetEndStrategySqlDispatcher dispatcher, String eventTable, Boolean retainCohortCovariates) {
+    return dispatcher.getStrategySql(this, eventTable, retainCohortCovariates);
   }  
 }

src/main/java/org/ohdsi/circe/cohortdefinition/DateOffsetStrategy.java

@@ -38,7 +38,7 @@ public enum DateField {
   public int offset = 0;
 
   @Override
-  public String accept(IGetEndStrategySqlDispatcher dispatcher, String eventTable) {
-    return dispatcher.getStrategySql(this, eventTable);
+  public String accept(IGetEndStrategySqlDispatcher dispatcher, String eventTable, Boolean retainCohortCovariates) {
+    return dispatcher.getStrategySql(this, eventTable, retainCohortCovariates);
   }
 }

src/main/java/org/ohdsi/circe/cohortdefinition/Death.java

@@ -19,7 +19,16 @@
 package org.ohdsi.circe.cohortdefinition;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+
+import org.apache.commons.lang3.StringUtils;
 import org.ohdsi.circe.cohortdefinition.builders.BuilderOptions;
+import org.ohdsi.circe.cohortdefinition.builders.ColumnFieldData;
+import org.ohdsi.circe.cohortdefinition.builders.ColumnFieldDataType;
 import org.ohdsi.circe.vocabulary.Concept;
 
 /**
@@ -53,4 +62,92 @@ public String accept(IGetCriteriaSqlDispatcher dispatcher, BuilderOptions option
     return dispatcher.getCriteriaSql(this, options);
   }
 
+  @Override
+  public List<ColumnFieldData> getSelectedField(Boolean retainCohortCovariates) {
+      List<ColumnFieldData> selectCols = new ArrayList<>();
+
+      if (retainCohortCovariates) {
+          if (occurrenceStartDate != null) {
+              selectCols.add(new ColumnFieldData("death_date", ColumnFieldDataType.DATE));
+          }
+
+          if (deathType != null && deathType.length > 0) {
+              selectCols.add(new ColumnFieldData("death_type_concept_id", ColumnFieldDataType.INTEGER));
+          }
+
+          if (deathSourceConcept != null) {
+              selectCols.add(new ColumnFieldData("cause_concept_id", ColumnFieldDataType.INTEGER));
+          }
+      }
+      return selectCols;
+  }
+
+  @Override
+  public String embedWindowedCriteriaQuery(String query, Map<String, ColumnFieldData> mapDistinctField) {
+      List<String> selectCols = new ArrayList<>();
+      List<String> groupCols = new ArrayList<>();
+      for (Entry<String, ColumnFieldData> entry : mapDistinctField.entrySet()) {
+          if (entry.getKey().equals("death_date") && occurrenceStartDate != null) {
+              selectCols.add(", cc.death_date");
+              groupCols.add(", cc.death_date");
+          } else if (entry.getKey().equals("death_type_concept_id") && deathType != null && deathType.length > 0) {
+              selectCols.add(", cc.death_type_concept_id");
+              groupCols.add(", cc.death_type_concept_id");
+          } else if (entry.getKey().equals("cause_concept_id") && deathSourceConcept != null) {
+              selectCols.add(", cc.cause_concept_id");
+              groupCols.add(", cc.cause_concept_id");
+          } else {
+              selectCols.add(", CAST(null as " + entry.getValue().getDataType().getType() + ") " + entry.getKey());
+          }
+      }
+
+      query = StringUtils.replace(query, "@additionColumnscc", StringUtils.join(selectCols, ""));
+      query = StringUtils.replace(query, "@additionGroupColumnscc", StringUtils.join(groupCols, ""));
+      return query;
+  }
+
+  @Override
+  public String embedWindowedCriteriaQueryP(String query) {
+      ArrayList<String> selectColsA = new ArrayList<>();
+
+      if (occurrenceStartDate != null) {
+          selectColsA.add(", A.death_date");
+      }
+
+      if (deathType != null && deathType.length > 0) {
+          selectColsA.add(", A.death_type_concept_id");
+      }
+
+      if (deathSourceConcept != null) {
+          selectColsA.add(", A.cause_concept_id");
+      }
+
+      query = StringUtils.replace(query, "@p.additionColumns", StringUtils.join(selectColsA, ""));
+      return query;
+  }
+
+  @Override
+  public String embedWrapCriteriaQuery(String query, List<String> selectColsPE, BuilderOptions options) {
+      ArrayList<String> selectCols = new ArrayList<>();
+
+      if(!options.isPrimaryCriteria()){
+        if (occurrenceStartDate != null) {
+            selectCols.add(", Q.death_date");
+            selectColsPE.add(", PE.death_date");
+        }
+
+        if (deathType != null && deathType.length > 0) {
+            selectCols.add(", Q.death_type_concept_id");
+            selectColsPE.add(", PE.death_type_concept_id");
+        }
+
+        if (deathSourceConcept != null) {
+            selectCols.add(", Q.cause_concept_id");
+            selectColsPE.add(", PE.cause_concept_id");
+        }
+      }
+
+      query = StringUtils.replace(query, "@QAdditionalColumnsInclusionN", StringUtils.join(selectCols, ""));
+      return query;
+  }
 }