Cannot merge tabix files with different sequence names

src/htsjdk/java/htsjdk/samtools/BAMIndexMerger.java

@@ -8,6 +8,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * Merges BAM index files for (headerless) parts of a BAM file into a single
@@ -93,12 +94,15 @@ private static BAMIndexContent mergeBAMIndexContent(int referenceSequence,
    */
   public static BinningIndexContent.BinList mergeBins(List<BinningIndexContent.BinList> binLists, long[] offsets) {
     List<Bin> mergedBins = new ArrayList<>();
-    int maxBinNumber = binLists.stream().mapToInt(bl -> bl.maxBinNumber).max().orElse(0);
+    int maxBinNumber = binLists.stream().filter(Objects::nonNull).mapToInt(bl -> bl.maxBinNumber).max().orElse(0);
     int commonNonNullBins = 0;
     for (int i = 0; i <= maxBinNumber; i++) {
       List<Bin> nonNullBins = new ArrayList<>();
       for (int j = 0; j < binLists.size(); j++) {
         BinningIndexContent.BinList binList = binLists.get(j);
+        if (binList == null) {
+          continue;
+        }
         Bin bin = VirtualShiftUtil.shift(binList.getBin(i), offsets[j]);
         if (bin != null) {
           nonNullBins.add(bin);
@@ -110,7 +114,7 @@ public static BinningIndexContent.BinList mergeBins(List<BinningIndexContent.Bin
       }
     }
     int numberOfNonNullBins =
-        binLists.stream().mapToInt(BinningIndexContent.BinList::getNumberOfNonNullBins).sum() - commonNonNullBins;
+        binLists.stream().filter(Objects::nonNull).mapToInt(BinningIndexContent.BinList::getNumberOfNonNullBins).sum() - commonNonNullBins;
     return new BinningIndexContent.BinList(mergedBins.toArray(new Bin[0]), numberOfNonNullBins);
   }
 
@@ -209,6 +213,9 @@ private static BAMIndexMetaData mergeMetaData(List<BAMIndexMetaData> metaDataLis
   public static LinearIndex mergeLinearIndexes(int referenceSequence, List<LinearIndex> linearIndexes, long[] offsets) {
     int maxIndex = -1;
     for (LinearIndex li : linearIndexes) {
+      if (li == null) {
+        continue;
+      }
       if (li.getIndexStart() != 0) {
         throw new IllegalArgumentException("Cannot merge linear indexes that don't all start at zero");
       }
@@ -223,6 +230,9 @@ public static LinearIndex mergeLinearIndexes(int referenceSequence, List<LinearI
     for (int i = 0; i < maxIndex; i++) {
       for (int liIndex = 0; liIndex < linearIndexes.size(); liIndex++) {
         LinearIndex li = linearIndexes.get(liIndex);
+        if (li == null) {
+          continue;
+        }
         long[] indexEntries = li.getIndexEntries();
         // Use the first linear index that has an index entry at position i.
         // There is no need to check later linear indexes, since their entries

src/htsjdk/java/htsjdk/samtools/TabixPlainText.java

@@ -31,6 +31,9 @@ public static File textIndexTbi(File tbi) throws IOException {
   public static void dump(PrintWriter pw, TabixIndex tbi) {
     BinningIndexContent[] binningIndexContents = TabixIndexMerger.getBinningIndexContents(tbi);
     for (BinningIndexContent content : binningIndexContents) {
+      if (content == null) {
+        continue;
+      }
       writeReference(pw, content);
     }
   }

src/htsjdk/java/htsjdk/tribble/index/tabix/AllRefsTabixIndexCreator.java

@@ -0,0 +1,205 @@
+/*
+ * The MIT License
+ *
+ * Copyright (c) 2014 The Broad Institute
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+package htsjdk.tribble.index.tabix;
+
+import htsjdk.samtools.BinningIndexBuilder;
+import htsjdk.samtools.BinningIndexContent;
+import htsjdk.samtools.Chunk;
+import htsjdk.samtools.SAMSequenceDictionary;
+import htsjdk.samtools.SAMSequenceRecord;
+import htsjdk.tribble.Feature;
+import htsjdk.tribble.index.Index;
+import htsjdk.tribble.index.IndexCreator;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+// This is a copy of htsjdk.tribble.index.tabix.TabixIndexCreator, except sequence
+// names are populated from the header, not from the ones that are seen. This
+// change is needed to support index merging.
+
+/**
+ * IndexCreator for Tabix.
+ * Features are expected to be 1-based, inclusive.
+ */
+public class AllRefsTabixIndexCreator implements IndexCreator {
+    private final TabixFormat formatSpec;
+    private final List<BinningIndexContent> indexContents = new ArrayList<BinningIndexContent>();
+    private final SAMSequenceDictionary sequenceDictionary;
+
+    private int currentReferenceIndex = -1;
+    private BinningIndexBuilder indexBuilder = null;
+    // A feature can't be added to the index until the next feature is added because the next feature
+    // defines the location of the end of the previous feature in the output file.
+    private TabixFeature previousFeature = null;
+
+
+    /**
+     * @param sequenceDictionary is not required, but if present all features added must refer to sequences in the
+     *                           dictionary.  It is used to optimize the memory needed to build the index.
+     */
+    public AllRefsTabixIndexCreator(final SAMSequenceDictionary sequenceDictionary,
+                                    final TabixFormat formatSpec) {
+        this.sequenceDictionary = sequenceDictionary;
+        this.formatSpec = formatSpec.clone();
+    }
+
+    @Override
+    public void addFeature(final Feature feature, final long filePosition) {
+        final String sequenceName = feature.getContig();
+        final int referenceIndex = sequenceDictionary.getSequenceIndex(sequenceName);
+        boolean advance = false;
+        if (currentReferenceIndex == -1) {
+            for (int i = 0; i < referenceIndex; i++) { // add nulls if not 0th referenceIndex
+                indexContents.add(null);
+            }
+            currentReferenceIndex = referenceIndex;
+            advance = true;
+        } else {
+            if (referenceIndex == currentReferenceIndex + 1) {
+                advance = true;
+            }
+            if (referenceIndex != currentReferenceIndex && referenceIndex != currentReferenceIndex + 1) {
+                throw new IllegalArgumentException("Sequence " + feature + " added out of order" + (" currentReferenceIndex: " + currentReferenceIndex + ", referenceIndex:" + referenceIndex));
+            }
+        }
+        final TabixFeature thisFeature = new TabixFeature(referenceIndex, feature.getStart(), feature.getEnd(), filePosition);
+        if (previousFeature != null) {
+            if (previousFeature.compareTo(thisFeature) > 0) {
+                throw new IllegalArgumentException(String.format("Features added out of order: previous (%s) > next (%s)",
+                        previousFeature, thisFeature));
+            }
+            finalizeFeature(filePosition);
+        }
+        previousFeature = thisFeature;
+        if (advance) {
+            advanceToReference(referenceIndex);
+        }
+    }
+
+    private void finalizeFeature(final long featureEndPosition) {
+        previousFeature.featureEndFilePosition = featureEndPosition;
+        if (previousFeature.featureStartFilePosition >= previousFeature.featureEndFilePosition) {
+            throw new IllegalArgumentException(String.format("Feature start position %d >= feature end position %d",
+                    previousFeature.featureStartFilePosition, previousFeature.featureEndFilePosition));
+        }
+        indexBuilder.processFeature(previousFeature);
+    }
+
+    private void advanceToReference(final int referenceIndex) {
+        if (indexBuilder != null) {
+            indexContents.add(indexBuilder.generateIndexContent());
+        }
+        // If sequence dictionary is provided, BinningIndexBuilder can reduce size of array it allocates.
+        final int sequenceLength;
+        if (sequenceDictionary != null) {
+            sequenceLength = sequenceDictionary.getSequence(referenceIndex).getSequenceLength();
+        } else {
+            sequenceLength = 0;
+        }
+        indexBuilder = new BinningIndexBuilder(referenceIndex, sequenceLength);
+        currentReferenceIndex = referenceIndex;
+    }
+
+    @Override
+    public Index finalizeIndex(final long finalFilePosition) {
+        if (previousFeature != null) {
+            finalizeFeature(finalFilePosition);
+        }
+        if (indexBuilder != null) {
+            indexContents.add(indexBuilder.generateIndexContent());
+        }
+        // Make this as big as the sequence dictionary, even if there is not content for every sequence,
+        // but truncate the sequence dictionary before its end if there are sequences in the sequence dictionary without
+        // any features.
+        final BinningIndexContent[] indices = indexContents.toArray(new BinningIndexContent[sequenceDictionary.size()]);
+        List<String> sequenceNames = sequenceDictionary.getSequences().stream().map(SAMSequenceRecord::getSequenceName).collect(Collectors.toList());
+        return new TabixIndex(formatSpec, sequenceNames, indices);
+    }
+
+
+    private static class TabixFeature implements BinningIndexBuilder.FeatureToBeIndexed, Comparable<TabixFeature> {
+        private final int referenceIndex;
+        private final int start;
+        private final int end;
+        private final long featureStartFilePosition;
+        // Position after this feature in the file.
+        private long featureEndFilePosition = -1;
+
+        private TabixFeature(final int referenceIndex, final int start, final int end, final long featureStartFilePosition) {
+            this.referenceIndex = referenceIndex;
+            this.start = start;
+            this.end = end;
+            this.featureStartFilePosition = featureStartFilePosition;
+        }
+
+        @Override
+        public int getStart() {
+            return start;
+        }
+
+        @Override
+        public int getEnd() {
+            return end;
+        }
+
+        /**
+         *
+         * @return null -- Let index builder compute this.
+         */
+        @Override
+        public Integer getIndexingBin() {
+            return null;
+        }
+
+        @Override
+        public Chunk getChunk() {
+            if (featureEndFilePosition == -1) {
+                throw new IllegalStateException("End position is not set");
+            }
+            return new Chunk(featureStartFilePosition, featureEndFilePosition);
+        }
+
+        @Override
+        public int compareTo(final TabixFeature other) {
+            final int ret = this.referenceIndex - other.referenceIndex;
+            if (ret != 0) return ret;
+            return this.start - other.start;
+        }
+
+        @Override
+        public String toString() {
+            return "TabixFeature{" +
+                    "referenceIndex=" + referenceIndex +
+                    ", start=" + start +
+                    ", end=" + end +
+                    ", featureStartFilePosition=" + featureStartFilePosition +
+                    ", featureEndFilePosition=" + featureEndFilePosition +
+                    '}';
+        }
+    }
+}

src/htsjdk/java/htsjdk/tribble/index/tabix/TabixIndexMerger.java

@@ -106,8 +106,8 @@ private static BinningIndexContent mergeBinningIndexContent(int referenceSequenc
     List<BinningIndexContent.BinList> binLists = new ArrayList<>();
     List<LinearIndex> linearIndexes = new ArrayList<>();
     for (BinningIndexContent binningIndexContent : binningIndexContentList) {
-      binLists.add(binningIndexContent.getBins());
-      linearIndexes.add(binningIndexContent.getLinearIndex());
+      binLists.add(binningIndexContent == null ? null : binningIndexContent.getBins());
+      linearIndexes.add(binningIndexContent == null ? null : binningIndexContent.getLinearIndex());
     }
     return new BinningIndexContent(referenceSequence, BAMIndexMerger.mergeBins(binLists, offsets), BAMIndexMerger.mergeLinearIndexes(referenceSequence, linearIndexes, offsets));
   }

src/main/java/org/disq_bio/disq/impl/formats/vcf/HeaderlessVcfOutputFormat.java

@@ -25,8 +25,8 @@
  */
 package org.disq_bio.disq.impl.formats.vcf;
 
+import htsjdk.tribble.index.IndexCreator;
 import htsjdk.tribble.index.tabix.TabixFormat;
-import htsjdk.tribble.index.tabix.TabixIndexCreator;
 import htsjdk.variant.variantcontext.VariantContext;
 import htsjdk.variant.variantcontext.writer.VCFWriterHelper;
 import htsjdk.variant.variantcontext.writer.VariantContextWriter;
@@ -63,7 +63,7 @@ static class VcfRecordWriter extends RecordWriter<Void, VariantContext> {
             "Cannot create tabix index for file that is not block compressed.");
       }
       OutputStream out = file.getFileSystem(conf).create(file);
-      TabixIndexCreator tabixIndexCreator;
+      IndexCreator tabixIndexCreator;
       if (tbiFile == null) {
         tabixIndexCreator = null;
       } else {

src/main/java/org/disq_bio/disq/impl/formats/vcf/StreamBasedTabixIndexCreator.java

@@ -29,9 +29,9 @@
 import htsjdk.samtools.SAMSequenceDictionary;
 import htsjdk.samtools.util.BlockCompressedOutputStream;
 import htsjdk.tribble.index.Index;
+import htsjdk.tribble.index.tabix.AllRefsTabixIndexCreator;
 import htsjdk.tribble.index.tabix.TabixFormat;
 import htsjdk.tribble.index.tabix.TabixIndex;
-import htsjdk.tribble.index.tabix.TabixIndexCreator;
 import htsjdk.tribble.index.tabix.TabixIndexMerger;
 import htsjdk.tribble.util.LittleEndianOutputStream;
 import java.io.IOException;
@@ -41,7 +41,7 @@
 
 // TODO: enhance htsjdk so that it's possible to write a tabix index to a stream (not just a path)
 // see IndexingVariantContextWriter, which calls writeBasedOnFeaturePath
-public class StreamBasedTabixIndexCreator extends TabixIndexCreator {
+public class StreamBasedTabixIndexCreator extends AllRefsTabixIndexCreator {
 
   static class StreamBasedTabixIndex extends TabixIndex {
     private final OutputStream out;

src/test/java/htsjdk/samtools/TbiEqualityChecker.java

@@ -89,7 +89,7 @@ private void assertEquals(
     Assert.assertEquals(
         "Number of non-null bins", bins1.getNumberOfNonNullBins(), bins2.getNumberOfNonNullBins());
     for (int i = 0; i <= bins1.maxBinNumber; i++) {
-      assertEquals(bins1.getBin(i), bins2.getBin(i), identical || i != bins1.maxBinNumber);
+      assertEquals(bins1.getBin(i), bins2.getBin(i), identical);
     }
   }