Bug fix for using a ReservoirDownsampler with a ReadsDownsamplingIterator. (#5594)

* Bug fix to allow a ReservoirDownsampler to be used within a ReadsDownsamplingIterator (a combination which previously didn't work correctly)

* As part of this fix, don't consider items in a ReservoirDownsampler to be finalized until signalEndOfInput() is called, and update client code to call signalEndOfInput() appropriately.

Resolves #4768

Author: cmnbroad

src/main/java/org/broadinstitute/hellbender/utils/downsampling/MutectDownsampler.java

@@ -94,6 +94,7 @@ private void finalizePendingReads() {
                 // the desired coverage, but if the region is decently mappable the shortfall will be minor.
                 final ReservoirDownsampler wellMappedDownsampler = new ReservoirDownsampler(maxCoverage, false);
                 pendingReads.stream().filter(read -> read.getMappingQuality() > SUSPICIOUS_MAPPING_QUALITY).forEach(wellMappedDownsampler::submit);
+                wellMappedDownsampler.signalEndOfInput();
                 final List<GATKRead> readsToFinalize = wellMappedDownsampler.consumeFinalizedItems();
                 if (stride > 1) {
                     Collections.sort(readsToFinalize, Comparator.comparingInt(GATKRead::getAssignedStart));

src/main/java/org/broadinstitute/hellbender/utils/downsampling/PositionalDownsampler.java

@@ -1,13 +1,15 @@
 package org.broadinstitute.hellbender.utils.downsampling;
 
 import htsjdk.samtools.SAMFileHeader;
+import org.broadinstitute.hellbender.exceptions.GATKException;
 import org.broadinstitute.hellbender.utils.Utils;
 import org.broadinstitute.hellbender.utils.read.GATKRead;
 import org.broadinstitute.hellbender.utils.read.ReadCoordinateComparator;
 import org.broadinstitute.hellbender.utils.read.ReadUtils;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 
 
 /**
@@ -71,16 +73,29 @@ private void handlePositionalChange( final GATKRead newRead ) {
         // Use ReadCoordinateComparator to determine whether we've moved to a new start position.
         // ReadCoordinateComparator will correctly distinguish between purely unmapped reads and unmapped reads that
         // are assigned a nominal position.
-        if ( previousRead != null && ReadCoordinateComparator.compareCoordinates(previousRead, newRead, header) != 0 ) {
-            if ( reservoir.hasFinalizedItems() ) {
-                finalizeReservoir();
+        if ( previousRead != null) {
+            final int cmpDiff = ReadCoordinateComparator.compareCoordinates(previousRead, newRead, header);
+            if (cmpDiff == 1) {
+                throw new IllegalStateException(
+                        String.format("Reads must be coordinate sorted (earlier %s later %s)", previousRead, newRead));
+            } else if (cmpDiff != 0) {
+                finalizeReservoir(true);
             }
         }
     }
 
-    private void finalizeReservoir() {
+    private void finalizeReservoir(final boolean expectFinalizedItems) {
+        // We can't consume finalized reads from the reservoir unless we first signal EOI.
+        // Once signalEndOfInput has been called and propagated to the ReservoirDownsampler, consumeFinalizedItems
+        // must be called on the ReservoirDownsampler before any new items can be submitted to it, to reset its
+        // state so it can be recycled/reused for the next downsampling position.
+        reservoir.signalEndOfInput();
+        if (expectFinalizedItems && ! reservoir.hasFinalizedItems() ) {
+            throw new GATKException.ShouldNeverReachHereException("Expected downsampled items to be present when none are");
+        }
         finalizedReads.addAll(reservoir.consumeFinalizedItems());
         reservoir.resetStats();
+        previousRead = null;
     }
 
     @Override
@@ -97,9 +112,11 @@ public List<GATKRead> consumeFinalizedItems() {
 
     @Override
     public boolean hasPendingItems() {
-        // The finalized items in the ReservoirDownsampler are pending items from the perspective of the
-        // enclosing PositionalDownsampler
-        return reservoir.hasFinalizedItems();
+        // The ReservoirDownsampler accumulates pending items until signalEndOfInput has been called, at which
+        // point all items that have survived downsampling become finalized. From the perspective of the enclosing
+        // PositionalDownsampler, both finalized items and pending items in the ReservoirDownsampler are considered
+        // pending.
+        return reservoir.hasFinalizedItems() || reservoir.hasPendingItems();
     }
 
     @Override
@@ -109,9 +126,11 @@ public GATKRead peekFinalized() {
 
     @Override
     public GATKRead peekPending() {
-        // The finalized items in the ReservoirDownsampler are pending items from the perspective of the
-        // enclosing PositionalDownsampler
-        return reservoir.peekFinalized();
+        // The ReservoirDownsampler accumulates pending items until signalEndOfInput has been called, at which
+        // point all items that have survived downsampling become finalized. From the perspective of the enclosing
+        // PositionalDownsampler, both finalized items and pending items in the ReservoirDownsampler are considered
+        // pending.
+        return Optional.ofNullable(reservoir.peekFinalized()).orElse(reservoir.peekPending());
     }
 
     @Override
@@ -121,7 +140,7 @@ public int size() {
 
     @Override
     public void signalEndOfInput() {
-        finalizeReservoir();
+        finalizeReservoir(false);
     }
 
     @Override
@@ -139,6 +158,7 @@ public boolean requiresCoordinateSortOrder() {
 
     @Override
     public void signalNoMoreReadsBefore( final GATKRead read ) {
+        Utils.nonNull(read, "Positional downsampler requires non-null reads");
         handlePositionalChange(read);
     }
 }

src/main/java/org/broadinstitute/hellbender/utils/downsampling/ReservoirDownsampler.java

@@ -51,6 +51,13 @@ public final class ReservoirDownsampler extends ReadsDownsampler {
      */
     private int totalReadsSeen;
 
+    /**
+     * In order to guarantee that all reads have equal probability of being discarded, we need to have consumed the
+     * entire input stream before any items can become finalized. All submitted items (that survive downsampling)
+     * remain pending until endOfInputStream is called, at which point they become finalized.
+     */
+    private boolean endOfInputStream;
+
 
     /**
      * Construct a ReservoirDownsampler
@@ -88,6 +95,9 @@ public ReservoirDownsampler(final int targetSampleSize ) {
     @Override
     public void submit ( final GATKRead newRead ) {
         Utils.nonNull(newRead, "newRead");
+        // Once the end of the input stream has been seen, consumeFinalizedItems or clearItems must be called to
+        // reset the state of the ReservoirDownsampler before more items can be submitted
+        Utils.validate(! endOfInputStream, "attempt to submit read after end of input stream has been signaled");
 
         // Only count reads that are actually eligible for discarding for the purposes of the reservoir downsampling algorithm
         totalReadsSeen++;
@@ -110,35 +120,50 @@ public void submit ( final GATKRead newRead ) {
 
     @Override
     public boolean hasFinalizedItems() {
-        return ! reservoir.isEmpty();
+        // All items in the reservoir are pending until endOfInputStream is seen, at which point all items become finalized
+        return endOfInputStream && !reservoir.isEmpty();
     }
 
     @Override
     public List<GATKRead> consumeFinalizedItems() {
+        // This method clears state (including the end of input stream flag) when called after
+        // end of input stream has been signaled, but has no side effects when called
+        // before end of input stream has been signaled (since in that case, the downsampling
+        // process is still ongoing and we shouldn't clear pending items).
+
         if (hasFinalizedItems()) {
             // pass reservoir by reference rather than make a copy, for speed
             final List<GATKRead> downsampledItems = reservoir;
             clearItems();
             return downsampledItems;
+        } else if ( ! endOfInputStream ) {
+            // Don't call clearItems() here, since endOfInputStream is false and therefore the
+            // downsampling process is still ongoing. We want to preserve existing pending items,
+            // and return an empty List without side effects.
+            return Collections.emptyList();
         } else {
-            // if there's nothing here, don't bother allocating a new list
+            // This is the case where endOfInputStream == true and our reservoir is empty. We return an empty
+            // list, but we also call clearItems() here for consistency with the case above where we have
+            // finalized items, so that in both cases we reset the endOfInputStream flag.
+            clearItems();
             return Collections.emptyList();
         }
     }
 
     @Override
     public boolean hasPendingItems() {
-        return false;
+        // All items in the reservoir are pending until endOfInputStream is seen, at which point all items become finalized
+        return !endOfInputStream && !reservoir.isEmpty();
     }
 
     @Override
     public GATKRead peekFinalized() {
-        return reservoir.isEmpty() ? null : reservoir.get(0);
+        return hasFinalizedItems() ? reservoir.get(0) : null;
     }
 
     @Override
     public GATKRead peekPending() {
-        return null;
+        return hasPendingItems() ? reservoir.get(0) : null;
     }
 
     @Override
@@ -148,7 +173,7 @@ public int size() {
 
     @Override
     public void signalEndOfInput() {
-        // NO-OP
+        endOfInputStream = true;
     }
 
     /**
@@ -164,6 +189,8 @@ public void clearItems() {
 
         // an internal stat used by the downsampling process, so not cleared by resetStats() below
         totalReadsSeen = 0;
+
+        endOfInputStream = false;
     }
 
     @Override

src/test/java/org/broadinstitute/hellbender/tools/HaplotypeCallerSparkIntegrationTest.java

@@ -145,7 +145,7 @@ public void testVCFModeIsConcordantWithGATK3_8Results() throws Exception {
      * Test that in VCF mode we're >= 99% concordant with GATK3.8 results
      * THIS TEST explodes with an exception because Allele-Specific annotations are not supported in vcf mode yet.
      * It's included to parallel the matching (also exploding) test for the non-spark HaplotypeCaller
-     * {@link org.broadinstitute.hellbender.tools.walkers.haplotypecaller.HaplotypeCallerIntegrationTest#testVCFModeIsConcordantWithGATK3_8ResultsAlleleSpecificAnnotations()}
+     * {@link org.broadinstitute.hellbender.tools.walkers.haplotypecaller.HaplotypeCallerIntegrationTest#testVCFModeIsConcordantWithGATK3_8ResultsAlleleSpecificAnnotations(String, String)}}
      */
     @Test(expectedExceptions = UserException.class)
     public void testVCFModeIsConcordantWithGATK3_8ResultsAlleleSpecificAnnotations() throws Exception {

src/test/java/org/broadinstitute/hellbender/utils/downsampling/ReadsDownsamplingIteratorUnitTest.java

@@ -9,6 +9,7 @@
 import org.testng.annotations.Test;
 
 import java.util.*;
+import java.util.stream.IntStream;
 
 public class ReadsDownsamplingIteratorUnitTest extends GATKBaseTest {
 
@@ -132,6 +133,22 @@ public void testRemoveThrows() {
         iter.remove();
     }
 
+    @Test
+    public void testReadsDownsamplingIteratorWithReservoirDownsampler() {
+        final int TOTAL_READ_COUNT = 100;
+        final int TARGET_COVERAGE = 45;
+        final List<GATKRead> reads = new ArrayList<>(TOTAL_READ_COUNT);
+        IntStream.range(1, TOTAL_READ_COUNT).forEach(i -> reads.add(readWithName(Integer.toString(i))));
+
+        final List<GATKRead> downsampledReads = new ArrayList<>();
+        final ReadsDownsamplingIterator downsamplingIter = new ReadsDownsamplingIterator(reads.iterator(), new ReservoirDownsampler(TARGET_COVERAGE));
+        for ( final GATKRead read : downsamplingIter ) {
+            downsampledReads.add(read);
+        }
+
+        Assert.assertEquals(downsampledReads.size(), TARGET_COVERAGE);
+    }
+
     private GATKRead readWithName( final String name ) {
         return ArtificialReadUtils.createArtificialRead(TextCigarCodec.decode("10M"), name);
     }

src/test/java/org/broadinstitute/hellbender/utils/downsampling/ReservoirDownsamplerUnitTest.java

@@ -61,29 +61,29 @@ public void testReservoirDownsampler(final ReservoirDownsamplerTest test ) {
         logger.warn("Running test: " + test);
 
         Utils.resetRandomGenerator();
-
         final ReadsDownsampler downsampler = new ReservoirDownsampler(test.reservoirSize);
-
         downsampler.submit(test.createReads());
 
+        // after submit, but before signalEndOfInput, all reads are pending, none are finalized
         if ( test.totalReads > 0 ) {
-            Assert.assertTrue(downsampler.hasFinalizedItems());
-            Assert.assertTrue(downsampler.peekFinalized() != null);
-            Assert.assertFalse(downsampler.hasPendingItems());
-            Assert.assertTrue(downsampler.peekPending() == null);
+            Assert.assertFalse(downsampler.hasFinalizedItems());
+            Assert.assertNull(downsampler.peekFinalized());
+            Assert.assertTrue(downsampler.hasPendingItems());
+            Assert.assertNotNull(downsampler.peekPending());
         }
         else {
             Assert.assertFalse(downsampler.hasFinalizedItems() || downsampler.hasPendingItems());
             Assert.assertTrue(downsampler.peekFinalized() == null && downsampler.peekPending() == null);
         }
 
+        // after signalEndOfInput, no reads are pending, all are finalized
         downsampler.signalEndOfInput();
 
         if ( test.totalReads > 0 ) {
             Assert.assertTrue(downsampler.hasFinalizedItems());
-            Assert.assertTrue(downsampler.peekFinalized() != null);
+            Assert.assertNotNull(downsampler.peekFinalized());
             Assert.assertFalse(downsampler.hasPendingItems());
-            Assert.assertTrue(downsampler.peekPending() == null);
+            Assert.assertNull(downsampler.peekPending());
         }
         else {
             Assert.assertFalse(downsampler.hasFinalizedItems() || downsampler.hasPendingItems());
@@ -102,6 +102,19 @@ public void testReservoirDownsampler(final ReservoirDownsamplerTest test ) {
 
         downsampler.resetStats();
         Assert.assertEquals(downsampler.getNumberOfDiscardedItems(), 0);
+
+        // use the same downsampling parameters, but this time consume the reads through a
+        // ReadsDownsamplingIterator, and validate that we get the same results as using downsampler directly
+        Utils.resetRandomGenerator();
+        final ReadsDownsampler downsamplerForIterator = new ReservoirDownsampler(test.reservoirSize);
+        final ReadsDownsamplingIterator downsamplingIterator = new ReadsDownsamplingIterator(
+                test.createReads().iterator(),
+                downsamplerForIterator);
+        final List<GATKRead> downsampledReadsFromIterator = new ArrayList<>(test.reservoirSize);
+        downsamplingIterator.forEach(downsampledReadsFromIterator::add);
+
+        Assert.assertEquals(downsamplerForIterator.getNumberOfDiscardedItems(), test.expectedNumDiscardedItems);
+        Assert.assertEquals(downsampledReadsFromIterator, downsampledReads);
     }
 
     @Test(expectedExceptions = IllegalArgumentException.class)
@@ -129,4 +142,61 @@ public void testNoNullSignalNoMoreReadsBefore() throws Exception {
         ReadsDownsampler rd = new ReservoirDownsampler(1, true);
         rd.signalNoMoreReadsBefore(null);
     }
+
+    // Calling consumeFinalizeItems() before end of input on a non-empty downsampler should
+    // return an empty List and not change the state of the downsampler:
+    @Test
+    public void testConsumeFinalizedItemsBeforeEndOfInput() {
+        final ReservoirDownsampler downsampler = new ReservoirDownsampler(10, true);
+        final GATKRead read = ArtificialReadUtils.createArtificialRead("100M");
+        downsampler.submit(read);
+
+        Assert.assertFalse(downsampler.hasFinalizedItems());
+        Assert.assertTrue(downsampler.hasPendingItems());
+
+        final List<GATKRead> returnedReads = downsampler.consumeFinalizedItems();
+        Assert.assertTrue(returnedReads.isEmpty());
+
+        // The downsampler should still have pending reads after the call to consumeFinalizedItems()
+        // (ie., the downsampling process should still be ongoing, and the state of the downsampler
+        // should have been preserved):
+        Assert.assertFalse(downsampler.hasFinalizedItems());
+        Assert.assertTrue(downsampler.hasPendingItems());
+    }
+
+    // Calling consumeFinalizedItems() after end of input on an empty downsampler should
+    // return an empty List and reset the end of input flag in the downsampler:
+    @Test
+    public void testConsumeFinalizedItemsAfterEndOfInputOnEmptyDownsampler() {
+        final ReservoirDownsampler downsampler = new ReservoirDownsampler(10, true);
+
+        Assert.assertFalse(downsampler.hasFinalizedItems());
+        Assert.assertFalse(downsampler.hasPendingItems());
+
+        downsampler.signalEndOfInput();
+
+        final List<GATKRead> returnedReads = downsampler.consumeFinalizedItems();
+        Assert.assertTrue(returnedReads.isEmpty());
+
+        Assert.assertFalse(downsampler.hasFinalizedItems());
+        Assert.assertFalse(downsampler.hasPendingItems());
+
+        // Submitting an additional read after both signalEndOfInput() and consumeFinalizedItems()
+        // should succeed (not throw), since the call to consumeFinalizedItems() should reset the
+        // end of input flag:
+        final GATKRead read = ArtificialReadUtils.createArtificialRead("100M");
+        downsampler.submit(read);
+
+        Assert.assertFalse(downsampler.hasFinalizedItems());
+        Assert.assertTrue(downsampler.hasPendingItems());
+    }
+
+    // Calling submit() directly after signalEndOfInput() should throw:
+    @Test(expectedExceptions = IllegalStateException.class)
+    public void testSubmitAfterEndOfInputThrows() {
+        final ReservoirDownsampler downsampler = new ReservoirDownsampler(10, true);
+        downsampler.signalEndOfInput();
+        final GATKRead read = ArtificialReadUtils.createArtificialRead("100M");
+        downsampler.submit(read);
+    }
 }