Allow VariantProcessingStep to supply custom merge code for scatter/gather

Author: bbimber

SequenceAnalysis/api-src/org/labkey/api/sequenceanalysis/pipeline/VariantProcessingStep.java

@@ -62,6 +62,11 @@ default void validateScatter(ScatterGatherMethod method, PipelineJob job) throws
         {
 
         }
+
+        default void performAdditionalMergeTasks(SequenceOutputHandler.JobContext ctx, PipelineJob job, TaskFileManager manager, ReferenceGenome genome, List<File> orderedScatterOutputs) throws PipelineJobException
+        {
+
+        }
     }
 
     public static interface MayRequirePrepareTask

SequenceAnalysis/src/org/labkey/sequenceanalysis/pipeline/ProcessVariantsHandler.java

@@ -35,6 +35,7 @@
 import org.labkey.api.sequenceanalysis.pipeline.SequenceAnalysisJobSupport;
 import org.labkey.api.sequenceanalysis.pipeline.SequenceOutputHandler;
 import org.labkey.api.sequenceanalysis.pipeline.SequencePipelineService;
+import org.labkey.api.sequenceanalysis.pipeline.TaskFileManager;
 import org.labkey.api.sequenceanalysis.pipeline.ToolParameterDescriptor;
 import org.labkey.api.sequenceanalysis.pipeline.VariantProcessingStep;
 import org.labkey.api.sequenceanalysis.run.SimpleScriptWrapper;
@@ -851,4 +852,18 @@ else if (AbstractGenomicsDBImportHandler.TILE_DB_FILETYPE.isType(input))
             throw new PipelineJobException("Unknown file type: " + input.getPath());
         }
     }
+
+    @Override
+    public void performAdditionalMergeTasks(JobContext ctx, PipelineJob job, TaskFileManager manager, ReferenceGenome genome, List<File> orderedScatterOutputs) throws PipelineJobException
+    {
+        List<PipelineStepCtx<VariantProcessingStep>> providers = SequencePipelineService.get().getSteps(job, VariantProcessingStep.class);
+        for (PipelineStepCtx<VariantProcessingStep> stepCtx : providers)
+        {
+            VariantProcessingStep step = stepCtx.getProvider().create(ctx);
+            if (step instanceof VariantProcessingStep.SupportsScatterGather)
+            {
+                ((VariantProcessingStep.SupportsScatterGather)step).performAdditionalMergeTasks(ctx, job, manager, genome, orderedScatterOutputs);
+            }
+        }
+    }
 }

SequenceAnalysis/src/org/labkey/sequenceanalysis/pipeline/VariantProcessingRemoteMergeTask.java

@@ -2,7 +2,6 @@
 
 import htsjdk.samtools.util.Interval;
 import org.apache.commons.lang3.StringUtils;
-import org.apache.logging.log4j.Logger;
 import org.jetbrains.annotations.NotNull;
 import org.labkey.api.pipeline.AbstractTaskFactory;
 import org.labkey.api.pipeline.AbstractTaskFactorySettings;
@@ -15,14 +14,12 @@
 import org.labkey.api.sequenceanalysis.SequenceOutputFile;
 import org.labkey.api.sequenceanalysis.pipeline.ReferenceGenome;
 import org.labkey.api.sequenceanalysis.pipeline.SequenceOutputHandler;
-import org.labkey.api.sequenceanalysis.run.AbstractDiscvrSeqWrapper;
+import org.labkey.api.sequenceanalysis.pipeline.VariantProcessingStep;
 import org.labkey.api.util.FileType;
-import org.labkey.api.writer.PrintWriters;
 import org.labkey.sequenceanalysis.run.variant.OutputVariantsStartingInIntervalsStep;
 
 import java.io.File;
 import java.io.IOException;
-import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
@@ -113,6 +110,7 @@ private VariantProcessingJob getPipelineJob()
         SequenceTaskHelper.logModuleVersions(getJob().getLogger());
         RecordedAction action = new RecordedAction(ACTION_NAME);
         TaskFileManagerImpl manager = new TaskFileManagerImpl(getPipelineJob(), _wd.getDir(), _wd);
+        JobContextImpl ctx = new JobContextImpl(getPipelineJob(), getPipelineJob().getSequenceSupport(), getPipelineJob().getParameterJson(), _wd.getDir(), new TaskFileManagerImpl(getPipelineJob(), _wd.getDir(), _wd), _wd);
 
         File finalOut;
         SequenceOutputHandler<SequenceOutputHandler.SequenceOutputProcessor> handler = getPipelineJob().getHandler();
@@ -122,7 +120,7 @@ private VariantProcessingJob getPipelineJob()
         }
         else
         {
-            finalOut = runDefaultVariantMerge(manager, action, handler);
+            finalOut = runDefaultVariantMerge(ctx, manager, action, handler);
         }
 
         Map<String, File> scatterOutputs = getPipelineJob().getScatterJobOutputs();
@@ -154,7 +152,7 @@ private VariantProcessingJob getPipelineJob()
         return new RecordedActionSet(action);
     }
 
-    private File runDefaultVariantMerge(TaskFileManagerImpl manager, RecordedAction action, SequenceOutputHandler<SequenceOutputHandler.SequenceOutputProcessor> handler) throws PipelineJobException
+    private File runDefaultVariantMerge(JobContextImpl ctx, TaskFileManagerImpl manager, RecordedAction action, SequenceOutputHandler<SequenceOutputHandler.SequenceOutputProcessor> handler) throws PipelineJobException
     {
         Map<String, List<Interval>> jobToIntervalMap = getPipelineJob().getJobToIntervalMap();
         getJob().setStatus(PipelineJob.TaskStatus.running, "Combining Per-Contig VCFs: " + jobToIntervalMap.size());
@@ -209,6 +207,15 @@ private File runDefaultVariantMerge(TaskFileManagerImpl manager, RecordedAction
             manager.addIntermediateFile(new File(vcf.getPath() + ".tbi"));
         }
 
+        Set<Integer> genomeIds = new HashSet<>();
+        getPipelineJob().getFiles().forEach(x -> genomeIds.add(x.getLibrary_id()));
+        if (genomeIds.size() != 1)
+        {
+            throw new PipelineJobException("Expected a single genome, found: " + StringUtils.join(genomeIds, ", "));
+        }
+
+        ReferenceGenome genome = getPipelineJob().getSequenceSupport().getCachedGenome(genomeIds.iterator().next());
+
         String basename = SequenceAnalysisService.get().getUnzippedBaseName(toConcat.get(0).getName());
         File combined = new File(getPipelineJob().getAnalysisDirectory(), basename + ".vcf.gz");
         File combinedIdx = new File(combined.getPath() + ".tbi");
@@ -223,18 +230,15 @@ private File runDefaultVariantMerge(TaskFileManagerImpl manager, RecordedAction
                 throw new PipelineJobException("Missing one of more VCFs: " + missing.stream().map(File::getPath).collect(Collectors.joining(",")));
             }
 
-            Set<Integer> genomeIds = new HashSet<>();
-            getPipelineJob().getFiles().forEach(x -> genomeIds.add(x.getLibrary_id()));
-            if (genomeIds.size() != 1)
-            {
-                throw new PipelineJobException("Expected a single genome, found: " + StringUtils.join(genomeIds, ", "));
-            }
-
-            ReferenceGenome genome = getPipelineJob().getSequenceSupport().getCachedGenome(genomeIds.iterator().next());
             combined = SequenceAnalysisService.get().combineVcfs(toConcat, combined, genome, getJob().getLogger(), true, null);
         }
         manager.addOutput(action, "Merged VCF", combined);
 
+        if (handler instanceof VariantProcessingStep.SupportsScatterGather)
+        {
+            ((VariantProcessingStep.SupportsScatterGather) handler).performAdditionalMergeTasks(ctx, getPipelineJob(), manager, genome, toConcat);
+        }
+
         return combined;
     }
 }

SequenceAnalysis/src/org/labkey/sequenceanalysis/run/variant/SplitVcfBySamplesStep.java

@@ -73,7 +73,7 @@ public Output processVariants(File inputVCF, File outputDirectory, ReferenceGeno
             if (!f.getName().equals(inputVCF.getName()) && f.getName().startsWith(basename) && SequenceUtil.FILETYPE.vcf.getFileType().isType(f))
             {
                 output.addOutput(f, "Subset VCF");
-                output.addSequenceOutput(f, "Subset VCF: " + f.getName(), "VCF", null, null, genome.getGenomeId(), null);
+                output.addSequenceOutput(f, "Subset VCF: " + f.getName(), "VCF File", null, null, genome.getGenomeId(), null);
             }
         }