Add more support for scatter/gather

Author: bbimber

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

@@ -68,7 +68,7 @@ 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
+        default void performAdditionalMergeTasks(SequenceOutputHandler.JobContext ctx, PipelineJob job, TaskFileManager manager, ReferenceGenome genome, List<File> orderedScatterOutputs, List<String> orderedJobDirs) throws PipelineJobException
         {
 
         }

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

@@ -872,14 +872,14 @@ else if (AbstractGenomicsDBImportHandler.TILE_DB_FILETYPE.isType(input))
     }
 
     @Override
-    public void performAdditionalMergeTasks(JobContext ctx, PipelineJob job, TaskFileManager manager, ReferenceGenome genome, List<File> orderedScatterOutputs) throws PipelineJobException
+    public void performAdditionalMergeTasks(JobContext ctx, PipelineJob job, TaskFileManager manager, ReferenceGenome genome, List<File> orderedScatterOutputs, List<String> orderedJobDirs) throws PipelineJobException
     {
         List<PipelineStepCtx<VariantProcessingStep>> providers = SequencePipelineService.get().getSteps(job, VariantProcessingStep.class);
         for (PipelineStepCtx<VariantProcessingStep> stepCtx : providers)
         {
             if (stepCtx.getProvider() instanceof VariantProcessingStep.SupportsScatterGather ssg)
             {
-                ssg.performAdditionalMergeTasks(ctx, job, manager, genome, orderedScatterOutputs);
+                ssg.performAdditionalMergeTasks(ctx, job, manager, genome, orderedScatterOutputs, orderedJobDirs);
             }
         }
     }

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

@@ -238,7 +238,7 @@ private File runDefaultVariantMerge(JobContextImpl ctx, TaskFileManagerImpl mana
         if (handler instanceof VariantProcessingStep.SupportsScatterGather)
         {
             ctx.getLogger().debug("Running additional merge tasks");
-            ((VariantProcessingStep.SupportsScatterGather) handler).performAdditionalMergeTasks(ctx, getPipelineJob(), manager, genome, toConcat);
+            ((VariantProcessingStep.SupportsScatterGather) handler).performAdditionalMergeTasks(ctx, getPipelineJob(), manager, genome, toConcat, new ArrayList<>(jobToIntervalMap.keySet()));
         }
 
         return combined;

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

@@ -3,34 +3,39 @@
 import htsjdk.samtools.util.Interval;
 import org.apache.logging.log4j.Logger;
 import org.jetbrains.annotations.Nullable;
+import org.labkey.api.pipeline.PipelineJob;
 import org.labkey.api.pipeline.PipelineJobException;
 import org.labkey.api.sequenceanalysis.SequenceAnalysisService;
+import org.labkey.api.sequenceanalysis.SequenceOutputFile;
 import org.labkey.api.sequenceanalysis.pipeline.AbstractVariantProcessingStepProvider;
 import org.labkey.api.sequenceanalysis.pipeline.CommandLineParam;
 import org.labkey.api.sequenceanalysis.pipeline.PipelineContext;
 import org.labkey.api.sequenceanalysis.pipeline.PipelineStep;
 import org.labkey.api.sequenceanalysis.pipeline.PipelineStepProvider;
 import org.labkey.api.sequenceanalysis.pipeline.ReferenceGenome;
+import org.labkey.api.sequenceanalysis.pipeline.SequenceOutputHandler;
+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.pipeline.VariantProcessingStepOutputImpl;
 import org.labkey.api.sequenceanalysis.run.AbstractCommandPipelineStep;
 import org.labkey.api.sequenceanalysis.run.AbstractDiscvrSeqWrapper;
+import org.labkey.sequenceanalysis.pipeline.SequenceJob;
 import org.labkey.sequenceanalysis.util.SequenceUtil;
 
 import java.io.File;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
-public class SplitVcfBySamplesStep extends AbstractCommandPipelineStep<SplitVcfBySamplesStep.Wrapper> implements VariantProcessingStep
+public class SplitVcfBySamplesStep extends AbstractCommandPipelineStep<SplitVcfBySamplesStep.Wrapper> implements VariantProcessingStep, VariantProcessingStep.SupportsScatterGather
 {
     public SplitVcfBySamplesStep(PipelineStepProvider<?> provider, PipelineContext ctx)
     {
         super(provider, ctx, new Wrapper(ctx.getLogger()));
     }
 
-    public static class Provider extends AbstractVariantProcessingStepProvider<SelectSamplesStep>
+    public static class Provider extends AbstractVariantProcessingStepProvider<SelectSamplesStep> implements SupportsScatterGather
     {
         public Provider()
         {
@@ -67,20 +72,65 @@ public Output processVariants(File inputVCF, File outputDirectory, ReferenceGeno
 
         output.addInput(inputVCF, "Input VCF");
 
+        return output;
+    }
+
+    private List<File> findProducedVcfs(File inputVCF, File outputDirectory)
+    {
+        List<File> ret = new ArrayList<>();
         String basename = SequenceAnalysisService.get().getUnzippedBaseName(inputVCF.getName());
         for (File f : outputDirectory.listFiles())
         {
             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 File", null, null, genome.getGenomeId(), null);
+                ret.add(f);
             }
         }
 
-        return output;
+        return ret;
     }
 
+    @Override
+    public void performAdditionalMergeTasks(SequenceOutputHandler.JobContext ctx, PipelineJob job, TaskFileManager manager, ReferenceGenome genome, List<File> orderedScatterOutputs, List<String> orderedJobDirs) throws PipelineJobException
+    {
+        job.getLogger().info("Merging additional track VCFs");
+        File inputVCF = ((SequenceJob)getPipelineCtx().getJob()).getInputFiles().get(0);
+        List<File> firstJobOutputs = findProducedVcfs(inputVCF, new File(ctx.getWorkingDirectory(), orderedJobDirs.get(0)));
+        for (File fn : firstJobOutputs)
+        {
+            List<File> toConcat = orderedJobDirs.stream().map(jobDir -> {
+                File f = new File(new File(getPipelineCtx().getWorkingDirectory(), jobDir), fn.getName());
+                if (!f.exists())
+                {
+                    throw new IllegalStateException("Missing file: " + f.getPath());
+                }
+
+                ctx.getFileManager().addIntermediateFile(f);
+                ctx.getFileManager().addIntermediateFile(new File(f.getPath() + ".tbi"));
 
+                return f;
+            }).toList();
+
+            String basename = SequenceAnalysisService.get().getUnzippedBaseName(toConcat.get(0).getName());
+            File combined = new File(ctx.getSourceDirectory(), basename + ".vcf.gz");
+            File combinedIdx = new File(combined.getPath() + ".tbi");
+            if (combinedIdx.exists())
+            {
+                job.getLogger().info("VCF exists, will not recreate: " + combined.getPath());
+            }
+            else
+            {
+                combined = SequenceAnalysisService.get().combineVcfs(toConcat, combined, genome, job.getLogger(), true, null);
+            }
+
+            SequenceOutputFile so = new SequenceOutputFile();
+            so.setName("Subset VCF: " + fn);
+            so.setFile(combined);
+            so.setCategory("VCF File");
+            so.setLibrary_id(genome.getGenomeId());
+            manager.addSequenceOutput(so);
+        }
+    }
 
     public static class Wrapper extends AbstractDiscvrSeqWrapper
     {