Add bases-per-slice threshold to bound slice memory for long reads

Previously the slice flush decision was based only on a record count
threshold (e.g. 100,000 records for the archive profile). For long-read
data (PacBio HiFi 10-20kb, ONT), a full slice buffers 1+ GB of quality
scores, which then OOMs the FQZComp quality-score encoder.

Add a second threshold on accumulated bases per slice, modelled on
htslib's dual-threshold rule: flush when either records-per-slice or
bases-per-slice is reached. Default bases-per-slice is readsPerSlice *
500, matching htslib. This keeps short-read behaviour unchanged (500bp
* 10000 = 5MB fits comfortably) while capping PacBio/ONT slices at
~3,000-5,000 records.

- CRAMEncodingStrategy: add basesPerSlice field + get/set, with a
  derived default when not explicitly set.
- SliceFactory.getUpdatedReferenceContext: now takes numberOfBases
  alongside numberOfSAMRecords; flushes the slice if either threshold
  is reached.
- ContainerFactory: accumulate bases (via SAMRecord.getReadLength)
  across the active slice; pass through and reset on slice emit.

Author: tfenne

src/main/java/htsjdk/samtools/cram/build/ContainerFactory.java

@@ -74,6 +74,7 @@ public final class ContainerFactory {
     private final List<SAMRecord> sliceSAMRecords;
 
     private long globalRecordCounter = 0;
+    private long accumulatedBases = 0;
     private int currentReferenceContextID = ReferenceContext.UNINITIALIZED_REFERENCE_ID;
 
     /**
@@ -112,11 +113,13 @@ public final Container getNextContainer(final SAMRecord samRecord, final long co
         final int updatedReferenceContextID = sliceFactory.getUpdatedReferenceContext(
                 currentReferenceContextID,
                 nextRecordIndex,
-                sliceSAMRecords.size());
+                sliceSAMRecords.size(),
+                accumulatedBases);
 
         if (shouldEmitSlice(updatedReferenceContextID)) {
             sliceFactory.createNewSliceEntry(currentReferenceContextID, sliceSAMRecords);
             sliceSAMRecords.clear();
+            accumulatedBases = 0;
 
             if (shouldEmitContainer(
                     currentReferenceContextID,
@@ -130,6 +133,7 @@ public final Container getNextContainer(final SAMRecord samRecord, final long co
         }
 
         sliceSAMRecords.add(samRecord);
+        accumulatedBases += samRecord.getReadLength();
         return container;
     }
 

src/main/java/htsjdk/samtools/cram/build/SliceFactory.java

@@ -51,6 +51,7 @@ public final class SliceFactory {
 
     private long sliceRecordCounter;
     private final int maxRecordsPerSlice;
+    private final long maxBasesPerSlice;
     private final int minimumSingleReferenceSliceThreshold;
     private final boolean coordinateSorted;
 
@@ -71,6 +72,7 @@ public SliceFactory(
         this.cramReferenceRegion = new CRAMReferenceRegion(cramReferenceSource, samFileHeader.getSequenceDictionary());
         minimumSingleReferenceSliceThreshold = encodingStrategy.getMinimumSingleReferenceSliceSize();
         maxRecordsPerSlice = this.encodingStrategy.getReadsPerSlice();
+        maxBasesPerSlice = this.encodingStrategy.getBasesPerSlice();
         this.coordinateSorted = samFileHeader.getSortOrder() == SAMFileHeader.SortOrder.coordinate;
         this.sliceRecordCounter = globalRecordCounter;
         cramRecordSliceEntries = new ArrayList<>(this.encodingStrategy.getSlicesPerContainer());
@@ -174,7 +176,12 @@ private final List<CRAMCompressionRecord> convertToCRAMRecords(final List<SAMRec
 
     /**
      * Decide if the current records should be flushed based on the current reference context, the reference context
-     * for the next record to be written, and the number of records seen so far.
+     * for the next record to be written, the number of records seen so far, and the accumulated number of bases.
+     *
+     * A slice is considered "full" when either {@link CRAMEncodingStrategy#getReadsPerSlice} records or
+     * {@link CRAMEncodingStrategy#getBasesPerSlice} bases have been accumulated -- whichever happens first.
+     * This matches htslib behaviour and prevents pathologically large slices for long-read data
+     * (PacBio HiFi, ONT) where 100K records could otherwise buffer >1 GB of quality scores.
      *
      * Slices with the Multiple Reference flag (-2) set as the sequence ID in the header may contain reads mapped to
      * multiple external references, including unmapped reads (placed on these references or unplaced), but multiple
@@ -188,13 +195,17 @@ private final List<CRAMCompressionRecord> convertToCRAMRecords(final List<SAMRec
      * reference-differencing. This latter scenario is recommended for unsorted or non-coordinate-sorted data.
      *
      * @param nextReferenceIndex reference index of the next record to be emitted
+     * @param numberOfSAMRecords number of SAM records accumulated so far for the current slice
+     * @param numberOfBases total bases accumulated so far for the current slice
      * @return ReferenceContext.UNINITIALIZED_REFERENCE_ID if a current slice should be flushed and
      * subsequent records should go into a new slice; otherwise the updated reference context.
      */
     public int getUpdatedReferenceContext(
             final int currentReferenceContext,
             final int nextReferenceIndex,
-            final int numberOfSAMRecords) {
+            final int numberOfSAMRecords,
+            final long numberOfBases) {
+        final boolean sliceFull = numberOfSAMRecords >= maxRecordsPerSlice || numberOfBases >= maxBasesPerSlice;
         switch (currentReferenceContext) {
             // uninitialized can go to: unmapped or mapped
             case ReferenceContext.UNINITIALIZED_REFERENCE_ID:
@@ -207,9 +218,9 @@ public int getUpdatedReferenceContext(
             case ReferenceContext.UNMAPPED_UNPLACED_ID:
                 if (nextReferenceIndex == currentReferenceContext) {
                     // still unmapped...
-                    return numberOfSAMRecords < maxRecordsPerSlice ?
-                            ReferenceContext.UNMAPPED_UNPLACED_ID :
-                            ReferenceContext.UNINITIALIZED_REFERENCE_ID;
+                    return sliceFull ?
+                            ReferenceContext.UNINITIALIZED_REFERENCE_ID :
+                            ReferenceContext.UNMAPPED_UNPLACED_ID;
                 } else if (coordinateSorted) {
                     // coordinate sorted, and we're going from unmapped to mapped ??
                     throw new CRAMException("Invalid coord-sorted input - unmapped records must be last");
@@ -218,7 +229,7 @@ public int getUpdatedReferenceContext(
                     // record into the same slice with the unmapped ones, since there is no index query
                     // concern since we're not coord sorted anyway (though there is no reference compression
                     // happening in this container).
-                    return numberOfSAMRecords >= maxRecordsPerSlice ?
+                    return sliceFull ?
                             ReferenceContext.UNINITIALIZED_REFERENCE_ID :
                             ReferenceContext.MULTIPLE_REFERENCE_ID;
                 }
@@ -233,7 +244,7 @@ public int getUpdatedReferenceContext(
                             ReferenceContext.UNINITIALIZED_REFERENCE_ID; // emit a small mutli-ref
                 } else {
                     // multi-ref, not coord sorted
-                    return numberOfSAMRecords >= maxRecordsPerSlice ?
+                    return sliceFull ?
                             ReferenceContext.UNINITIALIZED_REFERENCE_ID :
                             ReferenceContext.MULTIPLE_REFERENCE_ID;
                 }
@@ -243,7 +254,7 @@ public int getUpdatedReferenceContext(
                 // (currentReferenceContext is an actual reference index, not a sentinel).
                 if (nextReferenceIndex == currentReferenceContext) {
                     // still on the same reference contig
-                    return numberOfSAMRecords >= maxRecordsPerSlice ?
+                    return sliceFull ?
                             ReferenceContext.UNINITIALIZED_REFERENCE_ID :
                             nextReferenceIndex;
                 } else {

src/main/java/htsjdk/samtools/cram/structure/CRAMCompressionProfile.java

@@ -52,9 +52,6 @@ public enum CRAMCompressionProfile {
      * <p>This profile uses trial compression: multiple codecs are tried per block and the smallest
      * result wins. Additional candidates include BZIP2, the Range (arithmetic) coder, and GZIP.
      */
-    // Note: large slices increase tag dictionary size; CompressionHeader.internalWrite()
-    // uses a fixed 100KB buffer for the tag dictionary, which may need adjustment for very
-    // large slice sizes.
     ARCHIVE(CramVersions.CRAM_v3_1, 7, 100_000);
 
     private final CRAMVersion cramVersion;

src/main/java/htsjdk/samtools/cram/structure/CRAMEncodingStrategy.java

@@ -46,11 +46,20 @@
 public class CRAMEncodingStrategy {
     public static final int DEFAULT_MINIMUM_SINGLE_REFERENCE_SLICE_THRESHOLD = 1000;
     public static final int DEFAULT_READS_PER_SLICE = 10000;
+    /**
+     * Default ratio of bases-per-slice to reads-per-slice. Matches htslib's default rule
+     * {@code bases_per_slice = seqs_per_slice * 500}. A slice is flushed when either the
+     * record count or the accumulated base count is reached, preventing individual slices
+     * from growing pathologically large when input reads are long (PacBio HiFi, ONT).
+     */
+    public static final int DEFAULT_BASES_PER_READ = 500;
 
     private CRAMVersion cramVersion = CramVersions.CRAM_v3_1;
     private int gzipCompressionLevel = Defaults.COMPRESSION_LEVEL;
     private int minimumSingleReferenceSliceSize = DEFAULT_MINIMUM_SINGLE_REFERENCE_SLICE_THRESHOLD;
     private int readsPerSlice = DEFAULT_READS_PER_SLICE;
+    // 0 means "derive from readsPerSlice at query time". Explicit override via setBasesPerSlice.
+    private long basesPerSlice = 0;
     private int slicesPerContainer = 1;
 
     private EnumMap<DataSeries, CompressorDescriptor> compressorMap;
@@ -146,6 +155,32 @@ public int getMinimumSingleReferenceSliceSize() {
         return minimumSingleReferenceSliceSize;
     }
 
+    /**
+     * Set the maximum accumulated bases per slice. When the accumulated bases in a slice
+     * reaches this threshold, the slice is flushed even if {@link #getReadsPerSlice} has
+     * not been reached.  This prevents individual slices from growing pathologically large
+     * for long-read data (PacBio HiFi, ONT).
+     *
+     * <p>Setting a value of 0 reverts to the default ({@link #getReadsPerSlice} {@code *}
+     * {@link #DEFAULT_BASES_PER_READ}), matching htslib's rule.
+     *
+     * @param basesPerSlice maximum bases per slice, or 0 to use the default
+     * @return this strategy for chaining
+     */
+    public CRAMEncodingStrategy setBasesPerSlice(final long basesPerSlice) {
+        ValidationUtils.validateArg(basesPerSlice >= 0, "basesPerSlice must be >= 0");
+        this.basesPerSlice = basesPerSlice;
+        return this;
+    }
+
+    /**
+     * @return the bases-per-slice threshold.  If no explicit value was set, returns
+     * {@link #getReadsPerSlice} {@code *} {@link #DEFAULT_BASES_PER_READ} (matching htslib).
+     */
+    public long getBasesPerSlice() {
+        return basesPerSlice > 0 ? basesPerSlice : (long) readsPerSlice * DEFAULT_BASES_PER_READ;
+    }
+
     /**
      * Set the GZIP compression level used for data series compressed with GZIP.
      *

src/test/java/htsjdk/samtools/cram/build/SliceFactoryTest.java

@@ -134,7 +134,7 @@ private void testEmitSliceMultiSliceContainer() {
                 0L);
         sliceFactory.createNewSliceEntry(0, CRAMStructureTestHelper.createSAMRecordsMapped(10, 0));
         Assert.assertEquals(
-                sliceFactory.getUpdatedReferenceContext(0, 1, 1),
+                sliceFactory.getUpdatedReferenceContext(0, 1, 1, 0),
                 ReferenceContext.UNINITIALIZED_REFERENCE_ID
         );
     }
@@ -151,7 +151,7 @@ private void testEmitSliceCoordinateSortedPositive(
                 CRAMStructureTestHelper.SAM_FILE_HEADER,
                 0L);
         Assert.assertEquals(
-                sliceFactory.getUpdatedReferenceContext(currentReferenceContext, nextReferenceContext, nRecordsSeen),
+                sliceFactory.getUpdatedReferenceContext(currentReferenceContext, nextReferenceContext, nRecordsSeen, 0),
                 expectedUpdatedReferenceContext
         );
     }
@@ -185,7 +185,7 @@ private void testEmitSliceCoordinateSortedNegative(
                 CRAMStructureTestHelper.REFERENCE_SOURCE,
                 CRAMStructureTestHelper.SAM_FILE_HEADER,
                 0L);
-        sliceFactory.getUpdatedReferenceContext(currentReferenceContext, nextReferenceContext, nRecordsSeen);
+        sliceFactory.getUpdatedReferenceContext(currentReferenceContext, nextReferenceContext, nRecordsSeen, 0);
     }
 
     @DataProvider(name="emitSliceQuerynameSortedPositive")
@@ -236,7 +236,7 @@ private void testEmitSliceQuerynameSortedPositive(
                 querySortedSAMFileHeader,
                 0L);
         Assert.assertEquals(
-                sliceFactory.getUpdatedReferenceContext(currentReferenceContext, nextReferenceContext, nRecordsSeen),
+                sliceFactory.getUpdatedReferenceContext(currentReferenceContext, nextReferenceContext, nRecordsSeen, 0),
                 expectedUpdatedReferenceContext
         );
     }