解决了bam和block不对齐情况的处理

Co-authored-by: Copilot <copilot@github.com>

.gitignore

@@ -8,6 +8,7 @@
 /build
 build.sh
 run.sh
+/output
 
 # Compiled Object files
 *.slo

convert.sh

@@ -0,0 +1 @@
+time samtools view -h -@ 32 ~/mini.bam | samtools view -@ 32 -h -b -o ~/mini-samtools.bam

src/sort/phase_1.cpp

@@ -23,10 +23,10 @@ void phase1Pipeline() {
     /* set up*/
     Phase1PipelineArg phase1Arg;
     phase1Arg.numThread = nsgv::gSortArg.NUM_THREADS;
-    const size_t kReadBufSize = 1L * 1024 * 1024 * phase1Arg.numThread;  // 平均每线程1M缓冲区，累加起来，用来读入文件（BAM/SAM）（相对解压之后的缓冲区，大小可以忽略）
+    const size_t kReadBufSize = 4L * 1024 * 1024 * phase1Arg.numThread;  // 平均每线程4M缓冲区，累加起来，用来读入文件（BAM/SAM）（相对解压之后的缓冲区，大小可以忽略）
 
     phase1Arg.uncompressBufBytes = nsgv::gSortArg.MAX_MEM * 0.9; // 比最大内存参数小点
-    phase1Arg.threadBlocksWrap.Resize(phase1Arg.numThread); // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容
+    phase1Arg.threadUncompressWrap.Resize(phase1Arg.numThread); // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容
     phase1Arg.singleThreadMemBytes = kReadBufSize * BAM_COMPRESS_RATIAO;
     spdlog::info("max mem: {}, uncompress mem: {}, single thread mem: {}", nsgv::gSortArg.MAX_MEM, phase1Arg.uncompressBufBytes, phase1Arg.singleThreadMemBytes);
     for (int i = 0; i < phase1Arg.READ_BUF_NUM; ++i) {

src/sort/phase_1.h

@@ -24,47 +24,74 @@
 #define START_IDX(i, nt, nele) ((i) * (nele) / (nt))
 #define STOP_IDX(i, nt, nele) (((i) + 1) * (nele) / (nt))
 
+
+// 把线程解压需要的数据放到一个结构体里
+struct ThreadUncompressData {
+    DataBuffer blockBuf; // 解压的block放在这里
+    size_t blockNum = 0; // 解压的block数量
+    BamArr bamArr; // 解析后的bam数据放在这里
+    BamArr firstBam; // 连接上一个不完整的bam数据，当作当前block的第一个bam
+    DataBuffer lastBamBuf;  // 最后那个不完整的bam解析时用到的缓冲区
+    size_t memOffset = 0; // 这个线程解压的数据在全局解压数据中的偏移位置
+    size_t bamOffset = 0; // 这个线程解析的bam数据在全局解压数据中的偏移位置
+
+
+    void Resize(int vecInitSize) {
+        blockBuf.AllocMem(vecInitSize * SINGLE_BLOCK_SIZE);  // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容
+        bamArr.arr.reserve(vecInitSize);
+        firstBam.arr.reserve(1);
+        lastBamBuf.AllocMem(SINGLE_BLOCK_SIZE);
+    }
+
+    void Reset() {
+        blockNum = 0;
+        blockBuf.Clear();
+        bamArr.Clear();
+        firstBam.Clear();
+        lastBamBuf.Clear();
+        memOffset = 0;
+        bamOffset = 0;
+    }
+
+    size_t GetBlockNum() const { return blockNum; }
+    size_t GetBamNum() const { return bamArr.Size() + firstBam.Size(); }
+};
+
 // 第一阶段用到的数据结构
 // 第一阶段的解压、排序、归并、压缩
-struct UncompressBlocksWrap {
+struct ThreadUncompressWrap {
-    vector<ThreadBlockArr> threadBlocks;  // 每个thread一个，用来保存解压后的block数据
+    vector<ThreadUncompressData> threadUncompressDataArr; // 每个线程一个，用来保存解压数据和解析bam数据
 
-    vector<DataBuffer> threadBlockBuf; // 每个线程一个，用来保存解压后的block数据，和上边的threadBlocks重复了，暂时保留，后续优化
+    ThreadUncompressWrap() {}
-
+    ThreadUncompressWrap(int numThread) { Resize(numThread); }
-    UncompressBlocksWrap() {}
+    ThreadUncompressWrap(int numThread, int vecInitSize) { Resize(numThread, vecInitSize); }
-    UncompressBlocksWrap(int numThread) { Resize(numThread); }
-    UncompressBlocksWrap(int numThread, int vecInitSize) { Resize(numThread, vecInitSize); }
     void Resize(int numThread) { Resize(numThread, 128); }
     void Resize(int numThread, int vecInitSize) {
-        threadBlocks.resize(numThread);
+        threadUncompressDataArr.resize(numThread);
-        threadBlockBuf.resize(numThread);
-        for (int i = 0; i < numThread; ++i) { threadBlocks[i].blockArr.reserve(vecInitSize); }
         for (int i = 0; i < numThread; ++i) {
-            threadBlocks[i].blockArr.reserve(vecInitSize);
+            threadUncompressDataArr[i].Resize(vecInitSize);
-            threadBlockBuf[i].allocMem(vecInitSize * SINGLE_BLOCK_SIZE);  // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容
         }
     }
     void ResetBlockArr() {
-        for (int i = 0; i < threadBlocks.size(); ++i) {
+        for (int i = 0; i < threadUncompressDataArr.size(); ++i) {
-            threadBlocks[i].clear();
+            threadUncompressDataArr[i].Reset();
-            threadBlockBuf[i].clear();
         }
     }
 
     uint64_t GetTotalBlockNum() {
         uint64_t totalBlockNum = 0;
-        for (int i = 0; i < threadBlocks.size(); ++i) {
+        for (int i = 0; i < threadUncompressDataArr.size(); ++i) {
-            totalBlockNum += threadBlocks[i].curIdx;
+            totalBlockNum += threadUncompressDataArr[i].GetBlockNum();
         }
         return totalBlockNum;
     }
 
-    uint64_t GetHeapBlockNum() {
+    uint64_t GetTotalBamNum() {
-        uint64_t heapBlockNum = 0;
+        uint64_t totalBamNum = 0;
-        for (int i = 0; i < threadBlocks.size(); ++i) {
+        for (int i = 0; i < threadUncompressDataArr.size(); ++i) {
-            heapBlockNum += threadBlocks[i].blockHeap.size();
+            totalBamNum += threadUncompressDataArr[i].GetBamNum();
         }
-        return heapBlockNum;
+        return totalBamNum;
     }
 };
 
@@ -77,7 +104,6 @@ struct Phase1PipelineArg {
 
     // common parameters
     int numThread = 0; // 线程数
-    uint64_t bamNum = 0; // 解压后的bam数量
     uint64_t singleThreadMemBytes = 0; // 单线程开辟的内存字节上限
     uint64_t uncompressBufBytes = 0; // 总的解压缓冲区大小
     uint64_t startBlockId = 0; // 当前轮次起始block id
@@ -85,14 +111,24 @@ struct Phase1PipelineArg {
     // for read-uncompress-parse
     uint64_t readOrder = 0; // 读取文件轮次编号，与下边的uncompressOrder对应
     uint64_t uncompressOrder = 0; // 并行解压gz block, 包含排序（缓冲区满之后排序），以及合并之后的压缩
+    uint64_t memCopyOrder = 0; // 串行拷贝解压数据到uncompressData的轮次编号，与上边的uncompressOrder对应
 
     volatile int readFinish = 0;
     yarn::lock_t* readSig;
     yarn::lock_t* uncompressSig;
     ReadBuffer readData[READ_BUF_NUM];  // 用来读如数据，双缓冲
-    UncompressBlocksWrap threadBlocksWrap;        // 每个thread一个，用来保存解压后的block数据
+    ThreadUncompressWrap threadUncompressWrap;        // 每个thread一个，用来保存解压后的block数据
     UncompressBlockBuffer uncompressData;  // 所有线程共用一个，串行往这里添加解压后的block数据
+    BamArr allBams; // 所有线程共用一个，串行往这里添加解析后的bam数据
 
+    // 判断bam是否有效的阈值
+    int maxSeqLen = 0; // bam里seq的最大长度，初始值是int的最大值，后续会根据解压的bam数据更新这个值，作为判断bam是否合法的一个条件
+    int maxBamLen = 0; // bam的最大长度，初始值是int的最大值，后续会根据解压的bam数据更新这个值，作为判断bam是否合法的一个条件
+
+    uint64_t bamNum = 0; // 解压后的bam数量
+    uint64_t blockNum = 0; // 解压后的block数量
+
+    int zeroStartBlockNum = 0;
 
     // for merge-compress-write
     uint64_t compressOrder = 0;  // 排序后压缩，这个和下边的writeOder对应，跟上边的order不相关

src/sort/phase_1_uncompress.cpp

@@ -22,47 +22,114 @@
 #include "util/profiling.h"
 #include "util/yarn.h"
 
-/* 多线程解压 */
+int GetBamLen(uint8_t* dataAddr) {
-static void mtUncompressBlock(void* data, long idx, int tid) {
+    uint32_t bamLen = 0;
-    PROF_T_BEG(mem_copy);
+    memcpy(&bamLen, dataAddr, 4);
+    if (nsgv::gIsBigEndian)
+        ed_swap_4p(&bamLen);
+    return bamLen;
+}
 
-    Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
+void ParseBam(uint8_t* dataAddr, OneBam& bam) {
-    ReadBuffer & readData = p.readData[p.uncompressOrder % p.READ_BUF_NUM];
+    uint32_t bams = 0;
+    bam.bamLen = GetBamLen(dataAddr);
+    dataAddr += 4;
+    bam.tid = le_to_u32(dataAddr);
+    bam.pos = le_to_i32(dataAddr + 4);
+    uint32_t x2 = le_to_u32(dataAddr + 8);
+    bam.qnameLen = x2 & 0xff;
+}
 
-    auto& blockArr = p.threadBlocksWrap.threadBlocks[tid];
+// 解析一个bam并放入bamArr
-    auto& blockItem = blockArr.add();
+void ParseAddBam(uint8_t* dataAddr, BamArr& bamArr) {
-    uint8_t* block = readData.startAddrArr[idx];
+    OneBam& bam = bamArr.Add();
+    ParseBam(dataAddr, bam);
+}
 
-    size_t dlen = SINGLE_BLOCK_SIZE;  // 65535
+// 返回解析bam的个数
-    int block_length = unpackInt16(&block[16]) + 1;
+size_t ParseAddAllBams(uint8_t* dataAddr, size_t startOffset, size_t endOffset, BamArr& bamArr, size_t* nextBamStartPtr = nullptr,
-    uint32_t crc = le_to_u32(block + block_length - 8);
+                       size_t* lastPosPtr = nullptr) {
-    int ret = bgzfUncompress(blockItem.data, &dlen, (Bytef*)block + BLOCK_HEADER_LENGTH, block_length - BLOCK_HEADER_LENGTH, crc);
+    size_t nextBamStart = startOffset;
-    if (ret != 0) {
+    size_t lastPos = 0;
-        spdlog::error("uncompress error, block id: {}, len: {}, ret: {}", idx, block_length, ret);
+    uint32_t bamLen = 0;
-        exit(0);
+    uint32_t bams = 0;
+
+    while (nextBamStart + 4 <= endOffset) {
+        uint8_t* curAddr = dataAddr + nextBamStart;
+        memcpy(&bamLen, curAddr, 4);
+        if (nsgv::gIsBigEndian)
+            ed_swap_4p(&bamLen);
+        nextBamStart += 4 + bamLen;
+        if (nextBamStart == endOffset) {               // 刚好解析到最后，说明这个block的内容都是完整bam
+            lastPos = endOffset;                       // 继续解析当前的block
+        } else if (nextBamStart > endOffset) {         // 当前bam不完整，不能解析
+            lastPos = nextBamStart - (4 + bamLen);  // 记录最后一个不完整的bam起始位置
+            break;
         }
-    blockItem.blockId = idx + p.startBlockId;
+        OneBam& bam = bamArr.Add();
-    blockItem.blockLen = dlen;
+        bam.bamLen = bamLen;
-    blockArr.blockHeap.push({blockItem.blockId, blockArr.curIdx - 1});  // 解压完成后，将block的id和在block数组里的索引加入堆中，方便后续排序和合并
+        curAddr += 4;
-
+        bam.tid = le_to_u32(curAddr);
-#if 0
+        bam.pos = le_to_i32(curAddr + 4);
-    // 放入全局缓冲区
+        uint32_t x2 = le_to_u32(curAddr + 8);
-    // spdlog::info("top id: {}, block id: {}", blockArr.blockHeap.top().blockId, p.uncompressData.nextBlockId);
+        bam.qnameLen = x2 & 0xff;
-    while (blockArr.blockHeap.top().blockId == p.uncompressData.nextBlockId) {
+        ++bams;
-        auto& top = blockArr.blockHeap.top();
-        // auto& topBlock = blockArr.blockArr[top.blockArrIdx];
-        // memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize, topBlock.data, topBlock.blockLen);
-        // p.uncompressData.startAddrArr.push_back(p.uncompressData.dataBuf + p.uncompressData.usedBufSize);
-        // p.uncompressData.usedBufSize += topBlock.blockLen;
-        // p.bamNum += topBlock.bamNum;
-
-        blockArr.blockHeap.pop();
-        p.uncompressData.nextBlockId += 1;
-        p.uncompressData.blockNum += 1;
     }
-#endif
+    if (nextBamStart < endOffset) {
+        lastPos = nextBamStart;  // 最后一个不完整的bam，连长度都不够数据解析
+    }
+    if (nextBamStartPtr) *nextBamStartPtr = nextBamStart;
+    if (lastPosPtr) *lastPosPtr = lastPos;
+    return bams;
+}
 
-    PROF_T_END(tid, mem_copy);
+// 检查当前内存对应的bam是否能正确解析
+static bool isValidBam(uint8_t* dataAddr, uint32_t &bamLen, int maxBamLen, int maxSeqLen) {
+    memcpy(&bamLen, dataAddr, 4);
+    if (nsgv::gIsBigEndian)
+        ed_swap_4p(&bamLen);
+    int32_t tid = 0;
+    int64_t pos = 0;
+    int qnameLen = 0;
+    uint16_t flag = 0;
+    uint8_t* x = dataAddr + 4;
+    uint32_t x2 = le_to_u32(x + 8);
+    qnameLen = x2 & 0xff;
+    if (qnameLen >= bamLen) {
+        return false;
+    }
+    int32_t seqLen = le_to_u32(x + 16);
+    if (seqLen >= bamLen) {
+        return false;
+    }
+    tid = le_to_u32(x);
+    pos = le_to_i32(x + 4);
+    uint32_t x3 = le_to_u32(x + 12);
+    flag = x3 >> 16;
+
+    if ((flag & BAM_FUNMAP) && tid == -1 && pos == -1) {
+        return true;
+    }
+    int nref = sam_hdr_nref(nsgv::gInHdr.header);
+    if (tid < 0 || tid >= nref) {
+        // 非法 tid
+        return false;
+    }
+    hts_pos_t ref_len = sam_hdr_tid2len(nsgv::gInHdr.header, tid);
+    if (pos < 0 || pos >= ref_len) {
+        // 非法 pos：0 ≤ pos < ref_len
+        return false;
+    }
+    // uint32_t n_cigar = x3 & 0xffff;
+    if (seqLen * 10 < maxSeqLen || maxSeqLen * 10 < seqLen) {  // 猜测条件可以仔细考虑下
+        //spdlog::info("invalid bam(seqlen), bamlen: {}, seqLen: {}, maxSeqLen: {}, maxSeqLen: {}", bamLen, seqLen, maxSeqLen, maxSeqLen);
+        //return false;
+    }
+    if (bamLen * 10 < maxBamLen || maxBamLen * 10 < bamLen) {  // 猜测条件可以仔细考虑下
+        //spdlog::info("invalid bam(bamlen), bamlen: {}, maxBamLen: {}, seqLen: {}, qnameLen: {}, n_cigar: {}", bamLen, maxBamLen, seqLen, qnameLen, n_cigar);
+        // return false;
+    }
+    return true;
 }
 
 // 多线程解压，静态分配任务，此时用idx代替tid，multi-thread uncompress bam blocks
@@ -76,12 +143,15 @@ static void mtUncompressBlockBatch(void* data, long idx, int tid) {
     int startIdx = START_IDX(idx, p.numThread, readData.startAddrArr.size());
     int stopIdx = STOP_IDX(idx, p.numThread, readData.startAddrArr.size());
 
-    auto &blockBuf = p.threadBlocksWrap.threadBlockBuf[tid];
+    auto &blockBuf = p.threadUncompressWrap.threadUncompressDataArr[tid].blockBuf;
+    auto &bamArr = p.threadUncompressWrap.threadUncompressDataArr[tid].bamArr;
 
+    // 开辟足够的内存
     if (stopIdx - startIdx > blockBuf.maxLen / SINGLE_BLOCK_SIZE) {
-        blockBuf.reAllocMem((stopIdx - startIdx) * SINGLE_BLOCK_SIZE);
+        blockBuf.ReAllocMem((stopIdx - startIdx) * SINGLE_BLOCK_SIZE);
     }
 
+    // 解压block
     for (int i = startIdx; i < stopIdx; ++i) {
         uint8_t* block = readData.startAddrArr[i];
         size_t dlen = SINGLE_BLOCK_SIZE;  // 65535
@@ -94,71 +164,161 @@ static void mtUncompressBlockBatch(void* data, long idx, int tid) {
         }
         blockBuf.curLen += dlen;
     }
+
+    // 用readPos来表示第一个bam起始位置，默认是ngs的bam，如果是三代bam，应该不需要计算这个了，因为三代bam很长
+    blockBuf.readPos = 0;
+    for (int i = 0; i < SINGLE_BLOCK_SIZE; ++i) {
+        uint32_t nextBamStart = i;  // 这个要注意，每次应该要计算一下
+        uint32_t bamLen = 0;
+        uint8_t nextPos = 0; // 是否检查下一个可能的位置
+        while (nextBamStart + 4 <= SINGLE_BLOCK_SIZE) {
+            if (isValidBam(blockBuf.data + nextBamStart, bamLen, p.maxBamLen, p.maxSeqLen)) {
+                nextBamStart += 4 + bamLen;
+            } else {
+                nextPos = 1;
+                break;
+            }
+        }
+        if (!nextPos) {
+            blockBuf.readPos = i;
+            break;
+        }
+    }
+    // 解析bam
+    ParseAddAllBams(blockBuf.data, blockBuf.readPos, blockBuf.curLen, bamArr, nullptr, &blockBuf.lastPos);
     PROF_T_END(tid, mem_copy);
 }
 
+// 处理相邻线程的block数据，可能有bam跨越这两个线程的block（GATK的bam）
+static void handleAdjacentThreadBlock(Phase1PipelineArg& p) {
+    auto& uncompressData = p.uncompressData;
+    auto& threadUncompressDataArr = p.threadUncompressWrap.threadUncompressDataArr;
+    size_t offset = 0;  // 当前线程对应的全局数据的起始偏移量
+    size_t bamOffset = p.allBams.Size(); // 当前线程解析的bam在全局数据中的偏移量
+
+    for (int tid = 0; tid < p.numThread; ++tid) {
+        threadUncompressDataArr[tid].memOffset = offset;
+        threadUncompressDataArr[tid].bamOffset = bamOffset;
+        auto& blockBuf = threadUncompressDataArr[tid].blockBuf;
+        auto& bamArr = threadUncompressDataArr[tid].bamArr;
+        auto& firstBam = threadUncompressDataArr[tid].firstBam;
+        auto& lastBamBuf = threadUncompressDataArr[tid].lastBamBuf;
+
+        bool hasLastData = false; // 上一个block里有不完整bam数据
+        int lastDataLen = 0; // 上一个block里不完整bam数据的长度
+        int leftDataLen = blockBuf.readPos;  // 本轮剩余的不完整的bam数据
+
+        if (tid == 0) { // 第一个线程
+            // 检查一下bam的定位是否正确
+            hasLastData = uncompressData.usedBufSize != uncompressData.lastEndPos;
+            if (hasLastData || blockBuf.readPos > 0) {
+                lastDataLen = uncompressData.usedBufSize - uncompressData.lastEndPos;
+                leftDataLen = blockBuf.readPos;  // 本轮剩余的不完整的bam数据
+                lastBamBuf.MemCopy(uncompressData.dataBuf + uncompressData.lastEndPos, lastDataLen);
+                lastBamBuf.MemCopy(blockBuf.data, leftDataLen);
+            }
+        } else {
+            hasLastData = threadUncompressDataArr[tid - 1].blockBuf.curLen != threadUncompressDataArr[tid - 1].blockBuf.lastPos;
+            if (hasLastData || blockBuf.readPos > 0) {  // 上一轮有遗留数据
+                lastDataLen = threadUncompressDataArr[tid - 1].blockBuf.curLen - threadUncompressDataArr[tid - 1].blockBuf.lastPos;
+                leftDataLen = blockBuf.readPos;  // 本轮剩余的不完整的bam数据
+                lastBamBuf.MemCopy(threadUncompressDataArr[tid - 1].blockBuf.data + threadUncompressDataArr[tid - 1].blockBuf.lastPos, lastDataLen);
+                lastBamBuf.MemCopy(blockBuf.data, leftDataLen);
+            }
+        }
+        if (hasLastData || blockBuf.readPos > 0) {
+            int bamLen = GetBamLen(lastBamBuf.data);
+            int claculatedBamLen = lastDataLen + leftDataLen - 4;
+            int additionDataLen = 0;
+            if (bamLen != claculatedBamLen) {         // 猜测错了
+                if (lastDataLen + leftDataLen < 4) {  // 不够解析bam长度
+                    lastBamBuf.MemCopy(blockBuf.data + leftDataLen, 4);
+                    leftDataLen += 4;
+                    bamLen = GetBamLen(lastBamBuf.data);  // 真正的长度
+                }
+                additionDataLen = bamLen + 4 - (lastDataLen + leftDataLen);  // 还缺多少数据能解析出完整的bam
+                if (additionDataLen > 0) {
+                    lastBamBuf.MemCopy(blockBuf.data + leftDataLen, additionDataLen);
+                    leftDataLen += additionDataLen;
+                }
+                // 重新解析
+                bamArr.Clear();
+                ParseAddAllBams(blockBuf.data, leftDataLen, blockBuf.curLen, bamArr, nullptr, &blockBuf.lastPos);
+                spdlog::error("bam len mismatch {}: {}, {}", tid, bamLen, claculatedBamLen);
+            }
+            ParseAddBam(lastBamBuf.data, firstBam);
+        }
+        offset += threadUncompressDataArr[tid].blockBuf.curLen;
+        bamOffset += threadUncompressDataArr[tid].bamArr.Size() + threadUncompressDataArr[tid].firstBam.Size();
+    }
+}
+
 static void mtMemCopy(void* data, long idx, int tid) {
     Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
     tid = idx; // 静态分配任务，此时用idx代替tid
-    uint64_t offset = 0;
+
-    for (int i = 0; i < tid; ++i) {
+    auto& threadUncompressDataArr = p.threadUncompressWrap.threadUncompressDataArr;  // 每个thread一个，用来保存解压后的block数据
-        offset += p.threadBlocksWrap.threadBlockBuf[i].curLen;
+    auto &uncompressData = p.uncompressData;              // 所有线程共用一个，串行往这里添加解压后的block数据
+
+    // 拷贝bam未解析数据到全局的uncompressData里
+    memcpy(uncompressData.dataBuf + uncompressData.usedBufSize + threadUncompressDataArr[tid].memOffset, threadUncompressDataArr[tid].blockBuf.data,
+           threadUncompressDataArr[tid].blockBuf.curLen);
+           
+    // 拷贝解析的bam到全局数据里
+    size_t i = 0;
+    for (; i < threadUncompressDataArr[tid].firstBam.Size(); ++i) {
+        p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].firstBam.arr[i];
+    }
+    for (size_t j = 0; j < threadUncompressDataArr[tid].bamArr.Size(); ++i, ++j) {
+        p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].bamArr.arr[j];
+    }
+
+    if (tid == p.numThread - 1) { // 最后一个线程，更新全局uncompressData的usedBufSize
+        uncompressData.usedBufSize += threadUncompressDataArr[tid].memOffset + threadUncompressDataArr[tid].blockBuf.curLen;
+        uncompressData.lastEndPos = uncompressData.usedBufSize - (threadUncompressDataArr[tid].blockBuf.curLen - threadUncompressDataArr[tid].blockBuf.lastPos);
+        p.allBams.curIdx += threadUncompressDataArr[tid].bamOffset + threadUncompressDataArr[tid].bamArr.Size() + threadUncompressDataArr[tid].firstBam.Size();
     }
-    memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize + offset, p.threadBlocksWrap.threadBlockBuf[tid].data,
-           p.threadBlocksWrap.threadBlockBuf[tid].curLen);
 }
 
 /* 将gz block进行解压，并进行线程内排序 */
 static void doPhase1Uncompress(Phase1PipelineArg& p, int finish = 0) {
     PROF_G_BEG(uncompress);
     uint64_t blockNum = p.readData[p.uncompressOrder % p.READ_BUF_NUM].startAddrArr.size();
-    // kt_for(p.numThread, mtUncompressBlock, &p, blockNum);
+    p.blockNum += blockNum;
+
     kt_for(p.numThread, mtUncompressBlockBatch, &p, p.numThread);
-    // 串行拷贝所有blocks
+    PROF_G_END(uncompress);
+
+    // 并行拷贝所有blocks
     PROF_G_BEG(mem_copy);
-#if 1
+    handleAdjacentThreadBlock(p);
+    p.allBams.Add(p.threadUncompressWrap.GetTotalBamNum());
+    p.bamNum += p.threadUncompressWrap.GetTotalBamNum();
     kt_for(p.numThread, mtMemCopy, &p, p.numThread);
-#else
-    for (int i = 0; i < p.numThread; ++i) {
-        memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize, p.threadBlocksWrap.threadBlockBuf[i].data, p.threadBlocksWrap.threadBlockBuf[i].curLen);
-        // p.uncompressData.startAddrArr.push_back(p.uncompressData.dataBuf + p.uncompressData.usedBufSize);
-        p.uncompressData.usedBufSize += p.threadBlocksWrap.threadBlockBuf[i].curLen;
-    }
-#endif
-
-#if 0
-    for (int i = 0; i < 1; ++i) {
-        auto& blockArr = p.threadBlocksWrap.threadBlocks[i];
-        for (int j = 0; j < blockArr.curIdx; ++j) {
-            auto& blockItem = blockArr.blockArr[j];
-            memcpy(p.uncompressData.dataBuf + p.uncompressData.usedBufSize, blockItem.data, blockItem.blockLen);
-            p.uncompressData.startAddrArr.push_back(p.uncompressData.dataBuf + p.uncompressData.usedBufSize);
-            p.uncompressData.usedBufSize += blockItem.blockLen;
-            p.uncompressData.blockNum += 1;
-        }
-    }
-#endif
-
     PROF_G_END(mem_copy);
 
-    p.startBlockId += blockNum;
+    PROF_G_BEG(parse_block);
+    PROF_G_END(parse_block);
 
     if (true) { // 缓冲区满了
-        spdlog::info("blocks num: {}, left: {}, uncompressed: {}", p.threadBlocksWrap.GetTotalBlockNum(), p.threadBlocksWrap.GetHeapBlockNum(),
+        spdlog::info("blocks num: {}, uncompressed: {}, bam num: {}, all bam num: {}, zero start blocks: {}", blockNum, p.blockNum,
-                     p.uncompressData.blockNum);
+                     p.threadUncompressWrap.GetTotalBamNum(), p.bamNum, p.zeroStartBlockNum);
-        p.uncompressData.Clear();
+        // p.uncompressData.Clear();
-        p.threadBlocksWrap.ResetBlockArr();
+        p.uncompressData.NextRound();
-        p.uncompressData.nextBlockId = p.startBlockId;
+        // spdlog::info("last data - 0: {}", p.uncompressData.usedBufSize - p.uncompressData.lastEndPos);
-        p.uncompressData.blockNum = 0;
+        p.threadUncompressWrap.ResetBlockArr();
+        //for (size_t i = 0; i < p.allBams.Size(); ++i) {
+        //    fprintf(gfp[0], "%d-%ld\n", p.allBams.arr[i].tid, p.allBams.arr[i].pos);
+        //}
+        p.allBams.Clear();
     }
 
-    PROF_G_END(uncompress);
 }
 
 /* phase1Uncompress step-2 解压线程 */
 void* phase1Uncompress(void* data) {
     Phase1PipelineArg& p = *(Phase1PipelineArg*)data;
-
+    int parseFirstBlock = 1;
     /* 2. do the work */
     while (true) {
         // previous dependency
@@ -171,7 +331,39 @@ void* phase1Uncompress(void* data) {
             }
             break;
         }
+        if (parseFirstBlock) {
+            parseFirstBlock = 0;
+            // 计算bam的平均长度，以及第一个block里的bam个数，用来指导后续的解压和排序
+            uint8_t* block = p.readData[p.uncompressOrder % p.READ_BUF_NUM].startAddrArr[0];
+            size_t dlen = SINGLE_BLOCK_SIZE;  // 65535
+            int block_length = unpackInt16(&block[16]) + 1;
+            uint32_t crc = le_to_u32(block + block_length - 8);
+            uint8_t oneBlock[SINGLE_BLOCK_SIZE];
+            int ret = bgzfUncompress(oneBlock, &dlen, (Bytef*)block + BLOCK_HEADER_LENGTH, block_length - BLOCK_HEADER_LENGTH, crc);
+            if (ret != 0) {
+                spdlog::error("First block uncompress error, len: {}, ret: {}", block_length, ret);
+                exit(0);
+            }
+            uint64_t nextBamStart = 0;  // 第一个block的起始bam位置
+            uint32_t bamLen = 0;
+            uint64_t allBamLen = 0;
+            uint64_t bamNum = 0;
+            /* 解析每个bam */
+            while (nextBamStart + 4 <= dlen) {
+                bamLen = GetBamLen(oneBlock + nextBamStart);
+                p.maxBamLen = p.maxBamLen < bamLen ? bamLen : p.maxBamLen;
+                uint8_t* x = oneBlock + nextBamStart + 4;
+                int32_t seqLen = le_to_u32(x + 16);
+                p.maxSeqLen = p.maxSeqLen < seqLen ? seqLen : p.maxSeqLen;
 
+                nextBamStart += 4 + bamLen;
+                allBamLen += bamLen;
+                ++bamNum;
+            }
+            p.uncompressData.avgBamSize = bamNum == 0 ? 0 : allBamLen / bamNum;
+            p.uncompressData.avgBamNumPerBlock = bamNum;
+            spdlog::info("avg bam size: {}, avg bam num per block: {}, max bam len: {}, max seq len: {}", p.uncompressData.avgBamSize, p.uncompressData.avgBamNumPerBlock, p.maxBamLen, p.maxSeqLen);
+        }
         doPhase1Uncompress(p);
 
         // update status

src/sort/sam_io.cpp

@@ -137,7 +137,7 @@ size_t readUncompressOneBlock(FILE *fpr, uint8_t *fBuf, DataBuffer *uDataPtr) {
     uint32_t crc = le_to_u32(fBuf + blockLen - 8);
     size_t newDataSize = uData.maxLen;
     while (uData.curLen + SINGLE_BLOCK_SIZE > newDataSize) newDataSize *= 2;
-    uData.reAllocMem(newDataSize);  // 需要重新开辟空间
+    uData.ReAllocMem(newDataSize);  // 需要重新开辟空间
     int ret = bgzfUncompress(&uData.data[uData.curLen], &dlen, (Bytef *)fBuf + BLOCK_HEADER_LENGTH,
                          blockLen - BLOCK_HEADER_LENGTH, crc);
     if (ret < 0) {
@@ -159,7 +159,7 @@ void parseSamHeader(FILE *fpr, HeaderBuf &hdrBuf) {
     int32_t i, nameLen, numNames = 0;
 
     header = sam_hdr_init();       // 初始化header
-    uData.allocMem(kMaxBlockSize); // 初始化解压数据的buffer
+    uData.AllocMem(kMaxBlockSize); // 初始化解压数据的buffer
     readUncompressOneBlock(fpr, fBuf, &uData); // 读取第一个gz block
 
     // 解析header

src/sort/sam_io.h

@@ -8,6 +8,7 @@
 struct DataBuffer {
     uint8_t *data;
     size_t readPos = 0; // 当前读取的位置
+    size_t lastPos = 0; // 最后一个bam的起始位置
     size_t curLen = 0;  // 当前使用的空间
     size_t maxLen = 0;  // 最大空间
 
@@ -25,19 +26,30 @@ struct DataBuffer {
         if (data)
             free(data);
     }
-    void allocMem(size_t memSize) {
+    void AllocMem(size_t memSize) {
         curLen = 0;
         maxLen = memSize;
         data = (uint8_t *)realloc(data, maxLen);
     }
 
-    void reAllocMem(size_t memSize) {
+    void ReAllocMem(size_t memSize) {
         if (memSize > maxLen) {
             maxLen = memSize;
             data = (uint8_t *)realloc(data, maxLen);
         }
     }
-    void clear() { curLen = 0; readPos = 0; }
+
+    void MemCopy(uint8_t *src, size_t len) {
+        ReAllocMem(curLen + len);
+        memcpy(&data[curLen], src, len);
+        curLen += len;
+    }
+
+    void Clear() {
+        curLen = 0;
+        readPos = 0;
+        lastPos = 0;
+    }
 };
 
 struct HeaderBuf {

src/sort/sort.cpp

@@ -661,10 +661,17 @@ static void samSortFirstPipe() {
 // 排序的入口函数，entry function
 int doSort() {
 
+#if 1
+    gfp[0] = fopen("f0.txt", "w");
+    gfp[1] = fopen("f1.txt", "w");
+    gfp[2] = fopen("f2.txt", "w");
+    gfp[3] = fopen("f3.txt", "w");
+#endif
+
 #if 1
     nsgv::gIsBigEndian = ed_is_big();
     // 第一轮排序，外排到多个中间文件
-    // bamSortFirstPipe();
+    //bamSortFirstPipe();
 
     phase1Pipeline();
 
@@ -701,6 +708,11 @@ int doSort() {
         if (bamp->l_data > 1000) {
             spdlog::info("large record len: {}", bamp->l_data);
         }
+        if (bam_num % 10000000 == 0) {
+            spdlog::info("bam num: {}, max bam len: {}", bam_num, max_bam_len);
+        }
+
+        // fprintf(gfp[0], "%d-%ld\n", bamp->core.tid, bamp->core.pos);
     }
     sam_close(inBamFp);
     spdlog::info("max record len: {}", max_bam_len);
@@ -708,6 +720,12 @@ int doSort() {
 
 #endif
 
+#if 1
+    fclose(gfp[0]);
+    fclose(gfp[1]);
+    fclose(gfp[2]);
+    fclose(gfp[3]);
+#endif
+
     return 0;
 }
-

src/sort/sort.h

@@ -19,10 +19,10 @@ using std::vector;
 struct ReadBuffer {
     uint8_t *dataBuf = nullptr;
     uint8_t *blockBuf = nullptr; // 用来保存上一轮没能完整读取的block
-    int readBufSize = 0; // 读入的buf大小
+    size_t readBufSize = 0; // 读入的buf大小
     vector<uint8_t *> startAddrArr; // 存放每个block的起始地址
     ReadBuffer() { }
-    ReadBuffer(int readBufSize_) {
+    ReadBuffer(size_t readBufSize_) {
         readBufSize = readBufSize_;
         dataBuf = (uint8_t *)malloc(readBufSize);
         blockBuf = (uint8_t *)malloc(SINGLE_BLOCK_SIZE);
@@ -31,7 +31,7 @@ struct ReadBuffer {
         if (dataBuf) free(dataBuf);
         if (blockBuf) free(blockBuf);
     }
-    void Resize(int readBufSize_) {
+    void Resize(size_t readBufSize_) {
         if (dataBuf) free(dataBuf);
         if (blockBuf) free(blockBuf);
         readBufSize = readBufSize_;
@@ -43,38 +43,72 @@ struct ReadBuffer {
 /* 用来串行保存解压后的block数据, 跟ReadBuffer差不多*/
 struct UncompressBlockBuffer {
     uint8_t *dataBuf = nullptr; // 用来保存解压后的block数据，串行往这里添加解压后的block数据
-    uint8_t *unCompleteBuf = nullptr; // 用来保存上一轮没能完整解析的bam数据（前一部分，需要后续拷贝到dataBuf里）
     uint64_t dataBufSize = 0; // 存放的解压之后的block的buf大小
     uint64_t usedBufSize = 0; // 已经使用的buf大小
-    uint64_t usedUnCompleteBufSize = 0; // 已经使用的unCompleteBuf大小
+    uint64_t lastEndPos = 0; // 上一轮计算完整bam结束位置，也是下一轮第一个bam（不完整）的起始数据
-    int unCompleteBufSize = 0; // unCompleteBuf里数据的字节数，即上一轮不完整的bam的前半部分大小
-    vector<uint8_t *> startAddrArr; // 存放每个bam的起始地址
-    uint64_t nextBlockId = 0; // 下一个需要放入的block id，按照顺序放入dataBuf里
 
-    uint64_t blockNum = 0; // 解压后的block数量
+    int avgBamSize = 0; // bam的平均长度，用来指导后续的解压和排序
+    int avgBamNumPerBlock = 0; // 每个block里bam的平均数量，用来指导后续的解压和排序
 
     UncompressBlockBuffer() { }
     UncompressBlockBuffer(uint64_t dataBufSize_) {
         dataBufSize = dataBufSize_;
         dataBuf = (uint8_t *)malloc(dataBufSize);
-        unCompleteBufSize = SINGLE_BLOCK_SIZE;
-        unCompleteBuf = (uint8_t*)malloc(SINGLE_BLOCK_SIZE);
     }
     ~UncompressBlockBuffer() {
         if (dataBuf) free(dataBuf);
-        if (unCompleteBuf) free(unCompleteBuf);
     }
     void Resize(uint64_t dataBufSize_) {
         if (dataBuf) free(dataBuf);
-        if (unCompleteBuf) free(unCompleteBuf);
         dataBufSize = dataBufSize_;
-        unCompleteBufSize = SINGLE_BLOCK_SIZE;
         dataBuf = (uint8_t *)malloc(dataBufSize);
-        unCompleteBuf = (uint8_t *)malloc(SINGLE_BLOCK_SIZE);
     }
-    void Clear() { usedBufSize = 0; usedUnCompleteBufSize = 0; startAddrArr.clear(); }
+    void Clear() {
+        usedBufSize = 0;
+        lastEndPos = 0;
+    }
+
+    void NextRound() {
+        usedBufSize -= lastEndPos;
+        memcpy(dataBuf, dataBuf + lastEndPos, usedBufSize);
+        lastEndPos = 0;
+    }
 };
 
+/* 对vector的一个包装，避免频繁内存分配和释放 */
+template <typename T>
+struct FastVector {
+    vector<T> arr;
+    size_t curIdx = 0;  // 当前已经添加的元素数量，也是下一个要添加的元素的索引
+    T& Add() {
+        if (curIdx < arr.size())
+            return arr[curIdx++];
+        else {
+#if 0
+            arr.resize((arr.size() + 1) << 1);
+            return arr[curIdx++];
+
+#else
+            arr.push_back(T());
+            curIdx++;
+            return arr.back();
+#endif
+        }
+    }
+    void Add(size_t num) {
+        if (curIdx + num > arr.size()) {
+            arr.resize(curIdx + num);
+        }
+    }
+    void ReAllocate(size_t num) {
+        if (num > arr.size()) {
+            arr.resize(num);
+        }
+    }
+    size_t Size() const { return curIdx; }
+    size_t Capacity() const { return arr.size(); }
+    void Clear() { curIdx = 0; }
+};
 
 /*  */
 
@@ -94,8 +128,8 @@ struct OneBam {
     uint16_t bamLen = 0;
     uint16_t qnameLen;    // 序列名字长度
     uint32_t offset = 0;  // 距离首地址的偏移量
-    uint32_t tid = 0;     // 比对到的染色体
+    int32_t tid = 0;     // 比对到的染色体
-    uint64_t pos = 0;     // mapping 位置
+    int64_t pos = 0;     // mapping 位置
     ThreadBlockArr* blockThread;
     // for test
     // bam1_t b;
@@ -145,8 +179,6 @@ struct ThreadBlockArr {
     vector<OneBlock> blockArr;  // 解压后的数据
     int curIdx = 0;             // 当前解压数据对应的vector的索引
     uint64_t bamNum = 0;        // 解压后的bam数量
-    // 最小堆
-    std::priority_queue<BlockIdIdx, std::vector<BlockIdIdx>, std::greater<BlockIdIdx>> blockHeap;
 
     OneBlock& add() {
         if (curIdx < blockArr.size())
@@ -172,7 +204,6 @@ struct ThreadBlockArr {
     void clear() {
         curIdx = 0;
         bamNum = 0;
-        blockHeap = {};
     }
 };
 
@@ -355,6 +386,10 @@ struct MergeSortData {
     }
 };
 
+typedef FastVector<OneBam> BamArr;
+typedef FastVector<OneBlock> BlockArr;
+
+
 /* 第一阶段的多线程流水线参数 */
 struct FirstPipeArg {
     static const int READ_BUF_NUM = 2; // 读入的buf数量

src/util/profiling.cpp

@@ -9,6 +9,8 @@
 uint64_t tprof[LIM_THREAD_PROF_TYPE][LIM_THREAD] = {0};
 uint64_t proc_freq = 1000;
 uint64_t gprof[LIM_GLOBAL_PROF_TYPE] = {0};
+
+FILE* gfp[4] = {NULL, NULL, NULL, NULL};
 #endif
 
 uint64_t realtimeMsec(void) {

src/util/profiling.h

@@ -2,6 +2,7 @@
 #include <stdint.h>
 #include <stdlib.h>
 #include <sys/time.h>
+#include <stdio.h>
 
 // #define SHOW_PERF
 
@@ -19,6 +20,8 @@ extern "C" {
 extern uint64_t proc_freq;
 extern uint64_t tprof[LIM_THREAD_PROF_TYPE][LIM_THREAD];
 extern uint64_t gprof[LIM_GLOBAL_PROF_TYPE];
+
+extern FILE* gfp[4];
 #endif
 
 #ifdef SHOW_PERF