[EGGO-30] Generate flattened, partitioned data.

Author: tomwhite

eggo/dag.py

@@ -336,6 +336,40 @@ def output(self):
                                           edition=self.edition))
 
 
+class ADAMFlattenPartitionTask(Task):
+
+    adam_command = Parameter()
+    allowed_file_formats = Parameter()
+    source_edition = 'flat'
+    edition = 'flat_locuspart'
+
+    def requires(self):
+        return ADAMBasicTask(adam_command=self.adam_command,
+                             allowed_file_formats=self.allowed_file_formats)
+
+    def run(self):
+        adam_cmd = ('{adam_home}/bin/adam-submit --master {spark_master_url} partition'
+                    ' -partition_strategy_file {partition_strategy_file}'
+                    ' {source} {target}').format(
+            adam_home=os.environ['ADAM_HOME'],
+            spark_master_url=os.environ['SPARK_MASTER_URL'],
+            partition_strategy_file='genotypes-flat-partition-strategy.json',
+            source=target_s3n_url(ToastConfig().config['name'],
+                                  edition=self.source_edition),
+            target=target_s3n_url(ToastConfig().config['name'],
+                                  edition=self.edition))
+        p = Popen(adam_cmd, shell=True)
+        p.wait()
+
+        if p.returncode == 0:
+            create_SUCCESS_file(target_s3_url(ToastConfig().config['name'],
+                                              edition=self.edition))
+
+    def output(self):
+        return S3FlagTarget(target_s3_url(ToastConfig().config['name'],
+                                          edition=self.edition))
+
+
 class VCF2ADAMTask(Task):
 
     def requires(self):
@@ -344,7 +378,9 @@ def requires(self):
         flat = ADAMFlattenTask(adam_command='vcf2adam',
                                allowed_file_formats=['vcf'])
         locuspart = ADAMPartitionTask(adam_command='vcf2adam',
-                               allowed_file_formats=['vcf'])
+                                      allowed_file_formats=['vcf'])
+        flat_locuspart = ADAMFlattenPartitionTask(adam_command='vcf2adam',
+                                                  allowed_file_formats=['vcf'])
         dependencies = [basic]
         for edition in ToastConfig().config['editions']:
             if edition == 'basic':
@@ -353,6 +389,8 @@ def requires(self):
                 dependencies.append(flat)
             elif edition == 'locuspart':
                 dependencies.append(locuspart)
+            elif edition == 'flat_locuspart':
+                dependencies.append(flat_locuspart)
         return dependencies
 
     def run(self):
@@ -371,6 +409,8 @@ def requires(self):
                                allowed_file_formats=['sam', 'bam'])
         locuspart = ADAMPartitionTask(adam_command='transform',
                                       allowed_file_formats=['sam', 'bam'])
+        flat_locuspart = ADAMFlattenPartitionTask(adam_command='transform',
+                                                  allowed_file_formats=['sam', 'bam'])
         dependencies = [basic]
         for edition in ToastConfig().config['editions']:
             if edition == 'basic':
@@ -379,5 +419,7 @@ def requires(self):
                 dependencies.append(flat)
             elif edition == 'locuspart':
                 dependencies.append(locuspart)
+            elif edition == 'flat_locuspart':
+                dependencies.append(flat_locuspart)
         return dependencies
 

genotypes-flat-partition-strategy.json

@@ -0,0 +1,4 @@
+[
+  { "type": "identity", "source": "variant__contig__contigName", "name": "chr" },
+  { "type": "range", "source": "variant__start", "name": "pos", "range": 1000000 }
+]

genotypes-partition-strategy.json

@@ -1,4 +1,4 @@
 [
   { "type": "identity", "source": "variant.contig.contigName", "name": "chr" },
-  { "type": "range", "source": "variant.start", "name": "pos", "range": 10000 }
+  { "type": "range", "source": "variant.start", "name": "pos", "range": 1000000 }
 ]

test/registry/test-genotypes.json

@@ -3,7 +3,7 @@
     "title": "Test 1000 Genomes Project VCF data",
     "target": "test/genotypes",
     "dag": "VCF2ADAMTask",
-    "editions": ["basic", "locuspart"],
+    "editions": ["basic", "flat", "locuspart", "flat_locuspart"],
     "sources": [
         {"format": "vcf", "compression": true, "url": "https://github.com/bigdatagenomics/eggo/raw/master/test/resources/chr22.small.vcf.gz"}
     ]