Merge pull request #1618 from synthetichealth/limit-csv-lines

Limit resources per CSV file

Author: jawalonoski

src/main/java/org/mitre/synthea/export/CSVConstants.java

@@ -27,26 +27,6 @@ public class CSVConstants {
   public static final String CLAIM_TRANSACTION_KEY = "claims_transactions";
   public static final String PATIENT_EXPENSE_KEY = "patient_expenses";
 
-  public static final String BASE_PATIENT_FILENAME = "patients.csv";
-  public static final String BASE_ALLERGY_FILENAME = "allergies.csv";
-  public static final String BASE_MEDICATION_FILENAME = "medications.csv";
-  public static final String BASE_CONDITION_FILENAME = "conditions.csv";
-  public static final String BASE_CAREPLAN_FILENAME = "careplans.csv";
-  public static final String BASE_OBSERVATION_FILENAME = "observations.csv";
-  public static final String BASE_PROCEDURE_FILENAME = "procedures.csv";
-  public static final String BASE_IMMUNIZATION_FILENAME = "immunizations.csv";
-  public static final String BASE_ENCOUNTER_FILENAME = "encounters.csv";
-  public static final String BASE_IMAGING_STUDY_FILENAME = "imaging_studies.csv";
-  public static final String BASE_DEVICE_FILENAME = "devices.csv";
-  public static final String BASE_SUPPLY_FILENAME = "supplies.csv";
-  public static final String BASE_ORGANIZATION_FILENAME = "organizations.csv";
-  public static final String BASE_PROVIDER_FILENAME = "providers.csv";
-  public static final String BASE_PAYER_FILENAME = "payers.csv";
-  public static final String BASE_PAYER_TRANSITION_FILENAME = "payer_transitions.csv";
-  public static final String BASE_CLAIM_FILENAME = "claims.csv";
-  public static final String BASE_CLAIM_TRANSACTION_FILENAME = "claims_transactions.csv";
-  public static final String BASE_PATIENT_EXPENSE_FILENAME = "patient_expenses.csv";
-
   public static final String PATIENT_HEADER_LINE =
       "Id,BIRTHDATE,DEATHDATE,SSN,DRIVERS,PASSPORT,"
       + "PREFIX,FIRST,MIDDLE,LAST,SUFFIX,MAIDEN,MARITAL,RACE,ETHNICITY,GENDER,BIRTHPLACE,"
@@ -130,35 +110,8 @@ public class CSVConstants {
       + "HEALTHCARE_EXPENSES,INSURANCE_COSTS,COVERED_COSTS"
       + NEWLINE;
 
-  public static final Map<String, String> BASE_FILENAME_MAP = initializeFilenameMap();
   public static final Map<String, String> HEADER_LINE_MAP = initializeHeaderMap();
 
-  private static Map<String, String> initializeFilenameMap() {
-    Map<String, String> map = new HashMap<>();
-
-    map.put(PATIENT_KEY, BASE_PATIENT_FILENAME);
-    map.put(ALLERGY_KEY, BASE_ALLERGY_FILENAME);
-    map.put(MEDICATION_KEY, BASE_MEDICATION_FILENAME);
-    map.put(CONDITION_KEY, BASE_CONDITION_FILENAME);
-    map.put(CAREPLAN_KEY, BASE_CAREPLAN_FILENAME);
-    map.put(OBSERVATION_KEY, BASE_OBSERVATION_FILENAME);
-    map.put(PROCEDURE_KEY, BASE_PROCEDURE_FILENAME);
-    map.put(IMMUNIZATION_KEY, BASE_IMMUNIZATION_FILENAME);
-    map.put(ENCOUNTER_KEY, BASE_ENCOUNTER_FILENAME);
-    map.put(IMAGING_STUDY_KEY, BASE_IMAGING_STUDY_FILENAME);
-    map.put(DEVICE_KEY, BASE_DEVICE_FILENAME);
-    map.put(SUPPLY_KEY, BASE_SUPPLY_FILENAME);
-    map.put(ORGANIZATION_KEY, BASE_ORGANIZATION_FILENAME);
-    map.put(PROVIDER_KEY, BASE_PROVIDER_FILENAME);
-    map.put(PAYER_KEY, BASE_PAYER_FILENAME);
-    map.put(PAYER_TRANSITION_KEY, BASE_PAYER_TRANSITION_FILENAME);
-    map.put(CLAIM_KEY, BASE_CLAIM_FILENAME);
-    map.put(CLAIM_TRANSACTION_KEY, BASE_CLAIM_TRANSACTION_FILENAME);
-    map.put(PATIENT_EXPENSE_KEY, BASE_PATIENT_EXPENSE_FILENAME);
-
-    return Collections.unmodifiableMap(map);
-  }
-
   private static Map<String, String> initializeHeaderMap() {
     Map<String, String> map = new HashMap<>();
 

src/main/java/org/mitre/synthea/export/CSVFileManager.java

@@ -4,7 +4,9 @@
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.OutputStreamWriter;
+import java.lang.NumberFormatException;
 import java.nio.charset.Charset;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Arrays;
 import java.util.Collections;
@@ -16,6 +18,7 @@
 import org.apache.commons.io.output.NullOutputStream;
 import org.mitre.synthea.export.CSVConstants;
 import org.mitre.synthea.helpers.Config;
+import org.mitre.synthea.helpers.SimpleCSV;
 
 public class CSVFileManager {
   /**
@@ -26,16 +29,11 @@ public class CSVFileManager {
   private Path outputDirectory;
   private List<String> includedFiles;
   private List<String> excludedFiles;
-  private Map<String, String> filenameMap = initializeFilenameMap();
+  private Map<String, String> filenameMap = new HashMap<>();
   private Map<String, OutputStreamWriter> writerMap = new HashMap<>();
-
-  private Map<String, String> initializeFilenameMap() {
-    HashMap<String, String> map = new HashMap<>();
-
-    map.putAll(CSVConstants.BASE_FILENAME_MAP);
-
-    return map;
-  }
+  private Map<String, Integer> resourceCountMap = new HashMap<>();
+  private int maxLinesPerFile;
+  private int fileNumberDigits;
 
   /**
    * "No-op" writer to use to prevent writing to excluded files.
@@ -50,6 +48,8 @@ private Map<String, String> initializeFilenameMap() {
    */
   public CSVFileManager() {
     initializeAppend();
+    initializeMaxLinesPerFile();
+    initializeFileNumberDigits();
     initializeOutputDirectory();
     initializeIncludedAndExcludedFiles();
   }
@@ -58,6 +58,24 @@ private void initializeAppend() {
     append = Config.getAsBoolean("exporter.csv.append_mode");
   }
 
+  private void initializeMaxLinesPerFile() {
+    try {
+      maxLinesPerFile = Config.getAsInteger("exporter.csv.max_lines_per_file", 0);
+    } catch (NumberFormatException ex) {
+      // if the property is present but not a numeric string
+      maxLinesPerFile = 0;
+    }
+  }
+
+  private void initializeFileNumberDigits() {
+    try {
+      fileNumberDigits = Config.getAsInteger("exporter.csv.file_number_digits", 1);
+    } catch (NumberFormatException ex) {
+      // if the property is present but not a numeric string
+      fileNumberDigits = 1;
+    }
+  }
+
   private void initializeOutputDirectory() {
     File output = Exporter.getOutputFolder("csv", null);
     output.mkdirs();
@@ -72,6 +90,24 @@ private void initializeOutputDirectory() {
     }
   }
 
+  private boolean multipleFilesPerResource() {
+    return maxLinesPerFile > 0;
+  }
+
+  private String filename(String resourceKey) {
+    return resourceKey + ".csv";
+  }
+
+  private String filename(String resourceKey, int fileNumber) {
+    String formattedNumber = String.valueOf(fileNumber);
+
+    if (fileNumberDigits > 1) {
+      formattedNumber = String.format("%0" + fileNumberDigits + "d", fileNumber);
+    }
+
+    return resourceKey + "-" + formattedNumber + ".csv";
+  }
+
   private void initializeIncludedAndExcludedFiles() {
     String includedFilesStr = Config.get("exporter.csv.included_files", "").trim();
     String excludedFilesStr = Config.get("exporter.csv.excluded_files", "").trim();
@@ -126,6 +162,25 @@ private static List<String> propStringToList(String fileListString) {
     return files;
   }
 
+  private int incrementResourceCount(String resourceKey) {
+    Integer resourceCount = resourceCountMap.get(resourceKey);
+
+    if (resourceCount == null) {
+      resourceCount = 0;
+    }
+
+    resourceCount++;
+    resourceCountMap.put(resourceKey, resourceCount);
+
+    return resourceCount;
+  }
+
+  private boolean resourceIsExcluded(String resourceKey) {
+    String baseFilename = filename(resourceKey);
+    return (!includedFiles.isEmpty() && !includedFiles.contains(baseFilename))
+      || excludedFiles.contains(baseFilename);
+  }
+
   /**
    * Helper method to instantiate, if necessary, and return the writer for the
    * resource type's CSV file. Returns a "no-op" writer for any excluded files.
@@ -134,20 +189,87 @@ private static List<String> propStringToList(String fileListString) {
    *
    * @return OutputStreamWriter for the given resource type's CSV file
    */
-  private OutputStreamWriter getResourceWriter(String resourceKey) throws IOException {
-    String baseFilename = CSVConstants.BASE_FILENAME_MAP.get(resourceKey);
-    boolean excluded = (!includedFiles.isEmpty() && !includedFiles.contains(baseFilename))
-        || excludedFiles.contains(baseFilename);
-    if (excluded) {
+  private OutputStreamWriter initializeResourceWriter(String resourceKey) throws IOException {
+    if (resourceIsExcluded(resourceKey)) {
       return NO_OP;
     }
 
-    String filename = filenameMap.get(resourceKey);
+    String filename = filename(resourceKey);
     File file = outputDirectory.resolve(filename).toFile();
     // file writing may fail if we tell it to append to a file that doesn't already exist
     boolean appendToThisFile = append && file.exists();
 
-    return new OutputStreamWriter(new FileOutputStream(file, appendToThisFile), charset);
+    OutputStreamWriter writer =
+        new OutputStreamWriter(new FileOutputStream(file, appendToThisFile), charset);
+    if (!append) {
+      writer.write(CSVConstants.HEADER_LINE_MAP.get(resourceKey));
+    }
+
+    return writer;
+  }
+
+  /**
+   * Helper method to instantiate, if necessary, and return the writer for the
+   * resource type's CSV file. Returns a "no-op" writer for any excluded files.
+   *
+   * @param resourceKey Key from CSVConstants for the resource type being written
+   *
+   * @return OutputStreamWriter for the given resource type's CSV file
+   */
+  private OutputStreamWriter initializeResourceWriter(String resourceKey, int resourceCount)
+      throws IOException {
+    if (resourceIsExcluded(resourceKey)) {
+      return NO_OP;
+    }
+
+    if (append && resourceCount == 1) {
+      resourceCount = getResourceCount(resourceKey) + 1;
+      resourceCountMap.put(resourceKey, resourceCount);
+    }
+
+    int fileNumber = (resourceCount - 1) / maxLinesPerFile  + 1;
+    String filename = filename(resourceKey, fileNumber);
+
+    File file = outputDirectory.resolve(filename).toFile();
+    // file writing may fail if we tell it to append to a file that doesn't already exist
+    boolean appendToThisFile = append && file.exists();
+
+    OutputStreamWriter writer =
+        new OutputStreamWriter(new FileOutputStream(file, appendToThisFile), charset);
+    if (!append || resourceCount % maxLinesPerFile == 1) {
+      writer.write(CSVConstants.HEADER_LINE_MAP.get(resourceKey));
+    }
+
+    return writer;
+  }
+
+  private int getResourceCount(String resourceKey) throws IOException {
+    int fileNumber = 1;
+
+    String currentFilename = filename(resourceKey, fileNumber);
+    File file = outputDirectory.resolve(currentFilename).toFile();
+
+    if (file.exists()) {
+      do {
+        fileNumber++;
+        currentFilename = filename(resourceKey, fileNumber);
+        file = outputDirectory.resolve(currentFilename).toFile();
+      } while ((file.exists()));
+
+      fileNumber--;
+    }
+
+    currentFilename = filename(resourceKey, fileNumber);
+    file = outputDirectory.resolve(currentFilename).toFile();
+
+    int resourceCount = (fileNumber - 1) * maxLinesPerFile;
+
+    if (file.exists()) {
+      String csvData = new String(Files.readAllBytes(file.toPath()));
+      resourceCount += SimpleCSV.parse(csvData).size();
+    }
+
+    return resourceCount;
   }
 
   /**
@@ -158,13 +280,35 @@ private OutputStreamWriter getResourceWriter(String resourceKey) throws IOExcept
    * @return OutputStreamWriter for the given resource type's CSV file
    */
   public OutputStreamWriter getWriter(String resourceKey) throws IOException {
+    if (multipleFilesPerResource()) {
+      return getWriterForMultipleFiles(resourceKey);
+    }
+
     OutputStreamWriter writer = writerMap.get(resourceKey);
     if (writer == null) {
-      writer = getResourceWriter(resourceKey);
+      writer = initializeResourceWriter(resourceKey);
       writerMap.put(resourceKey, writer);
-      if (!append) {
-        writer.write(CSVConstants.HEADER_LINE_MAP.get(resourceKey));
+    }
+
+    return writer;
+  }
+
+  private OutputStreamWriter getWriterForMultipleFiles(String resourceKey) throws IOException {
+    if (resourceIsExcluded(resourceKey)) {
+      return NO_OP;
+    }
+
+    int resourceCount = incrementResourceCount(resourceKey);
+
+    OutputStreamWriter writer = writerMap.get(resourceKey);
+
+    if (resourceCount % maxLinesPerFile == 1) {
+      if (writer != null) {
+        writer.flush();
       }
+
+      writer = initializeResourceWriter(resourceKey, resourceCount);
+      writerMap.put(resourceKey, writer);
     }
 
     return writer;
@@ -177,7 +321,7 @@ public OutputStreamWriter getWriter(String resourceKey) throws IOException {
    */
   public void flushWriter(String resourceKey) throws IOException {
     synchronized (resourceKey) {
-      OutputStreamWriter writer = getWriter(resourceKey);
+      OutputStreamWriter writer = writerMap.get(resourceKey);
       if (writer != null) {
         writer.flush();
       }

src/main/resources/synthea.properties

@@ -53,6 +53,12 @@ exporter.csv.folder_per_run = false
 # NOTE: the csv exporter does not actively delete files, so if Run 1 you included a file, then Run 2 you exclude that file, the version from Run 1 will still be present
 exporter.csv.included_files =
 exporter.csv.excluded_files = patient_expenses.csv
+# if set to a number, the output for resource types with more than this number
+# of resources will be split among multiple files
+exporter.csv.max_lines_per_file =
+# When using multiple files, the numbers in the filenames will be zero padded to
+# this many digits.
+exporter.csv.file_number_digits =
 
 exporter.cpcds.export = false
 exporter.cpcds.append_mode = false