From eb876d4fe20670c037e06144a760174dfb3d7efb Mon Sep 17 00:00:00 2001 From: zzh Date: Thu, 11 Jun 2026 23:04:59 +0800 Subject: [PATCH] =?UTF-8?q?=E8=A7=A3=E5=86=B3=E4=BA=86phase1=E9=98=B6?= =?UTF-8?q?=E6=AE=B5=E8=AE=A1=E7=AE=97offset=E7=9A=84bug=EF=BC=8C=E7=AC=AC?= =?UTF-8?q?=E4=BA=8C=E9=98=B6=E6=AE=B5=E5=88=9D=E6=AD=A5=E5=AE=8C=E6=88=90?= =?UTF-8?q?=EF=BC=8C=E4=BD=86=E6=98=AF=E8=BF=98=E6=9C=89bug?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/sort/phase_1.h | 31 +- src/sort/phase_1_compress.cpp | 22 +- src/sort/phase_1_read.cpp | 10 +- src/sort/phase_1_read.h | 4 +- src/sort/phase_1_sort.cpp | 84 ++++- src/sort/phase_1_sort.h | 7 +- src/sort/phase_1_uncompress.cpp | 89 +++++- src/sort/phase_1_write.cpp | 43 ++- src/sort/phase_2.cpp | 27 +- src/sort/phase_2.h | 525 +++++++++++++++++++++++++++++++- src/sort/phase_2_merge.cpp | 309 +++++++++++++++++++ src/sort/phase_2_merge.h | 12 + src/sort/phase_2_read.cpp | 410 +++++++++++++++++++++++++ src/sort/phase_2_read.h | 16 + src/sort/phase_2_write.cpp | 106 +++++++ src/sort/phase_2_write.h | 14 + src/sort/sam_io.cpp | 17 +- src/sort/sam_io.h | 27 +- src/sort/sort.cpp | 28 +- src/sort/sort.h | 75 ++++- src/sort/sort_args.h | 2 +- src/util/profiling.cpp | 5 + src/util/profiling.h | 8 +- src/util/yarn.h | 8 + 24 files changed, 1814 insertions(+), 65 deletions(-) diff --git a/src/sort/phase_1.h b/src/sort/phase_1.h index c087ce7..63efe1b 100644 --- a/src/sort/phase_1.h +++ b/src/sort/phase_1.h @@ -25,9 +25,6 @@ using std::string; -#define START_IDX(i, nt, nele) ((i) * (nele) / (nt)) -#define STOP_IDX(i, nt, nele) (((i) + 1) * (nele) / (nt)) - // 把线程解压需要的数据放到一个结构体里 struct ThreadUncompressData { @@ -36,8 +33,8 @@ struct ThreadUncompressData { BamArr bamArr; // 解析后的bam数据放在这里 BamArr firstBam; // 连接上一个不完整的bam数据，当作当前block的第一个bam DataBuffer lastBamBuf; // 最后那个不完整的bam解析时用到的缓冲区 - size_t memOffset = 0; // 这个线程解压的数据在全局解压数据中的偏移位置 - size_t bamOffset = 0; // 这个线程解析的bam数据在全局解压数据中的偏移位置 + int64_t memOffset = 0; // 这个线程解压的数据在全局解压数据中的偏移位置 + int64_t bamOffset = 0; // 这个线程解析的bam数据在全局解压数据中的偏移位置 void Resize(int vecInitSize) { @@ -113,11 +110,14 @@ struct BlockBams { struct MergeCompressData { vector blockDataArr; // 待压缩的数据 vector compressDataArr; // 压缩后的数据 + int curIdx = 0; void Resize(int blockNum) { blockDataArr.resize(blockNum); compressDataArr.resize(blockNum); } + + void Clear() { curIdx = 0; } }; /* 第一阶段的多线程流水线参数 */ @@ -132,9 +132,8 @@ struct Phase1PipelineArg { int numThread = 0; // 线程数 uint64_t singleThreadMemBytes = 0; // 单线程开辟的内存字节上限 uint64_t uncompressBufBytes = 0; // 总的解压缓冲区大小 - uint64_t startBlockId = 0; // 当前轮次起始block id uint64_t maxMemBytes = 0; // 最大可使用内存，参数 - int compressLevel = 1; // 压缩级别，参数 + int compressLevel = 1; // 中间文件压缩级别，参数 // for read-uncompress-parse uint64_t readOrder = 0; // 读取文件轮次编号，与下边的uncompressOrder对应 @@ -153,6 +152,7 @@ struct Phase1PipelineArg { UncompressBlockBuffer uncompressData; // 所有线程共用一个，串行往这里添加解压后的block数据 BamArr allBams; // 所有线程共用一个，串行往这里添加解析后的bam数据 + // vector lastBams; // 判断bam是否有效的阈值，暂时没用了 int maxSeqLen = 0; // bam里seq的最大长度，初始值是int的最大值，后续会根据解压的bam数据更新这个值，作为判断bam是否合法的一个条件 @@ -183,17 +183,26 @@ struct Phase1PipelineArg { string midFileNamePrefix; // 中间文件名前缀 FILE* midFilePtr = nullptr; // 中间文件指针 + // 统计用 + vector uncompressedDataSize; // nthread数组长度 + vector compressedDataSize; + Phase1PipelineArg() { readSig = yarn::NEW_LOCK(0); uncompressSig = yarn::NEW_LOCK(0); mergeSig = yarn::NEW_LOCK(0); compressSig = yarn::NEW_LOCK(0); lastRoundBuf.AllocMem(SINGLE_BLOCK_SIZE); + } - mergeBlocksThreshold = 1000; // 大概64M的解压缩后的数据 - for(int i=0; ioffset, bp->wholeBamLen); + // check bam + // CheckBam(p.uncompressData.dataBuf + bp->offset, bp->wholeBamLen); } - - // bgzfCompress(void* _dst, size_t* dlen, const void* src, size_t slen, int level) compressData.ReAllocMem(SINGLE_BLOCK_SIZE); // 压缩后的block数据不会超过单个block的大小 compressData.curLen = SINGLE_BLOCK_SIZE; // spdlog::info("block bytes: {}", blockData.curLen); @@ -49,13 +61,13 @@ static void doCompress(Phase1PipelineArg& p) { MergeCompressData& mergeCompressData = p.mergeCompressData[p.compressOrder % p.MERGE_BUF_NUM]; kt_for(p.numThread, mtCompressBlock, &p, mergeCompressData.blockDataArr.size()); //kt_for(1, mtCompressBlock, &p, mergeCompressData.blockDataArr.size()); + PROF_G_END(compress); compressBuf.Clear(); - for (int i=0; i < mergeCompressData.blockDataArr.size(); ++i) { + for (int i=0; i < mergeCompressData.blockDataArr.size(); ++i) { // 如果太慢，可以考虑并行拷贝，先计算偏移量，然后多线程拷贝 compressBuf.MemCopy(mergeCompressData.compressDataArr[i].data, mergeCompressData.compressDataArr[i].curLen); } - spdlog::info("compress bytes: {}", compressBuf.curLen); - PROF_G_END(compress); + // spdlog::info("compress bytes: {}", compressBuf.curLen); } /* phase1Compress step- 压缩线程 */ diff --git a/src/sort/phase_1_read.cpp b/src/sort/phase_1_read.cpp index 937d27a..5d4a3fb 100644 --- a/src/sort/phase_1_read.cpp +++ b/src/sort/phase_1_read.cpp @@ -23,8 +23,9 @@ /* 具体执行读取文件操作 */ /* 将bam文件内容读取到buf，解析buf中的gz block长度信息 */ -static size_t doPhase1ReadFile(Phase1PipelineArg& p, DataBuffer& halfBlock, FILE* fpr) { - ReadBuffer& readData = p.readData[p.readOrder % p.READ_BUF_NUM]; +// static size_t doPhase1ReadFile(Phase1PipelineArg& p, DataBuffer& halfBlock, FILE* fpr) { +size_t ReadBamFile(ReadBuffer& readData, DataBuffer& halfBlock, FILE* fpr) { + // ReadBuffer& readData = p.readData[p.readOrder % p.READ_BUF_NUM]; size_t readState = 0; size_t curReadPos = 0; int blockLen = 0; @@ -69,9 +70,6 @@ static size_t doPhase1ReadFile(Phase1PipelineArg& p, DataBuffer& halfBlock, FILE memcpy(halfBlock.data, &readData.dataBuf[curReadPos], halfBlock.readPos); // 将不完整的block拷贝到halfBlock } - // spdlog::info("block num-1: {}", readData.startAddrArr.size()); - // spdlog::info("read order: {}, max block len: {}", p.readOrder, maxBlockLen); - return readState; } @@ -90,7 +88,7 @@ void* phase1ReadFile(void* data) { yarn::DEPENDENCY_NOT_TO_BE(p.readSig, p.READ_BUF_NUM); PROF_G_BEG(read); - size_t readState = doPhase1ReadFile(p, halfBlock, fpr); + size_t readState = ReadBamFile(p.readData[p.readOrder % p.READ_BUF_NUM], halfBlock, fpr); PROF_G_END(read); if (readState == 0) { diff --git a/src/sort/phase_1_read.h b/src/sort/phase_1_read.h index cb2ebce..d839c35 100644 --- a/src/sort/phase_1_read.h +++ b/src/sort/phase_1_read.h @@ -10,4 +10,6 @@ #pragma once /* phase1ReadFile step-1 读取文件线程 */ -void* phase1ReadFile(void* data); \ No newline at end of file +void* phase1ReadFile(void* data); + +// size_t ReadBamFile(ReadBuffer& readData, DataBuffer& halfBlock, FILE* fpr); \ No newline at end of file diff --git a/src/sort/phase_1_sort.cpp b/src/sort/phase_1_sort.cpp index 5911ac3..312d94a 100644 --- a/src/sort/phase_1_sort.cpp +++ b/src/sort/phase_1_sort.cpp @@ -57,7 +57,7 @@ struct Phase1BamHeap { if (p == nullptr) { return -1; } - + // for (int i = 0; i < 4; ++i) { for (int i = 0; i < p->numThread; ++i) { size_t startIdx = START_IDX(i, p->numThread, p->allBams.Size()); size_t stopIdx = STOP_IDX(i, p->numThread, p->allBams.Size()); @@ -94,7 +94,7 @@ struct Phase1BamHeap { }; // 线程内排序 -static void mtInThreadSort(void* data, long idx, int tid) { +void mtInThreadSort(void* data, long idx, int tid) { Phase1PipelineArg& p = *(Phase1PipelineArg*)data; tid = idx; @@ -119,10 +119,34 @@ static void mtInThreadSort(void* data, long idx, int tid) { } } +void singleThreadSort(void* data) { + Phase1PipelineArg& p = *(Phase1PipelineArg*)data; + // 先按照坐标排序 + int startIdx = 0; + int stopIdx = p.allBams.Size(); + + if (nsgv::gSortArg.SORT_COORIDINATE) { + std::sort(p.allBams.arr.begin() + startIdx, p.allBams.arr.begin() + stopIdx, [](const OneBam& b1, const OneBam& b2) { + if (b1.tid == -1) + return false; + return b1.tid < b2.tid || (b1.tid == b2.tid && b1.pos < b2.pos) || (b1.tid == b2.tid && b1.pos == b2.pos && b1.qnameLen < b2.qnameLen); + }); + } else if (nsgv::gSortArg.QUERY_NAME_TYPE == nsmd::QueryNameType::PICARD) { + std::sort(p.allBams.arr.begin() + startIdx, p.allBams.arr.begin() + stopIdx, [&](const OneBam& b1, const OneBam& b2) { + int cmp = strncmp((char*)(p.uncompressData.dataBuf + b1.offset + OneBam::QnameOffset), + (char*)(p.uncompressData.dataBuf + b2.offset + OneBam::QnameOffset), std::min(b1.qnameLen, b2.qnameLen)); + if (cmp == 0) + return b1.qnameLen < b2.qnameLen; + return cmp < 0; + }); + } +} + void phase1Sort(void* data) { Phase1PipelineArg& p = *(Phase1PipelineArg*)data; PROF_G_BEG(sort); kt_for(p.numThread, mtInThreadSort, &p, p.numThread); + // kt_for(1, mtInThreadSort, &p, p.numThread); PROF_G_END(sort); // spdlog::info("b1:{}, b2:{}, b3:{}, b4:{}, b5:{}", p.allBams.arr[0].tid, p.allBams.arr[1].tid, p.allBams.arr[2].tid, p.allBams.arr[3].tid, @@ -133,6 +157,40 @@ void phase1Sort(void* data) { // p.allBams.arr[4].offset); } +template +void doMergeLastData(Phase1PipelineArg& p, Phase1BamHeap& heap) { + const OneBam* bam = nullptr; + + vector arr(p.allBams.Size()); + + //p.lastBams.resize(p.allBams.Size()); + size_t i = 0; + //uint64_t pos; + while ((bam = heap.Pop()) != nullptr) { + //pos += bam->tid; + arr[i++] = *bam; + //p.lastBams[i++] = bam; + } + //spdlog::info("bam num: {}, pos: {}", p.lastBams.size(), pos); + spdlog::info("bam size: {}-{}, buf size: {}-{}", p.allBams.Capacity(), arr.size(), p.uncompressData.dataBufSize, p.uncompressData.usedBufSize); + p.allBams.arr = arr; // 替换 +} + +void mergeLastData(void* data) { + Phase1PipelineArg& p = *(Phase1PipelineArg*)data; + + gUncompressDataBuf = p.uncompressData.dataBuf; // 主要用来比较bam的名字信息 + Phase1BamHeap posHeap; + Phase1BamHeap nameHeap; + if (nsgv::gSortArg.SORT_COORIDINATE) { + posHeap.Init(&p); + doMergeLastData(p, posHeap); + } else { + nameHeap.Init(&p); + doMergeLastData(p, nameHeap); + } +} + template bool doMergeSort(Phase1PipelineArg& p, Phase1BamHeap& heap) { auto &mergeCompressData = p.mergeCompressData[p.mergeOrder % p.MERGE_BUF_NUM]; @@ -153,8 +211,11 @@ bool doMergeSort(Phase1PipelineArg& p, Phase1BamHeap& heap) { mergeCompressData.blockDataArr[mergedBlockNum].Clear(); // 清理，为添加bam数据做准备 bamBytes = 0; } - // mergeCompressData.blockDataArr[mergedBlockNum].MemCopy(p.uncompressData.dataBuf + bam->offset, bam->wholeBamLen); +#if 0 + mergeCompressData.blockDataArr[mergedBlockNum].blockBuf.MemCopy(p.uncompressData.dataBuf + bam->offset, bam->wholeBamLen); +#else mergeCompressData.blockDataArr[mergedBlockNum].bamPtrArr.Add(bam); +#endif bamBytes += bam->wholeBamLen; // for test heap.Pop(); } @@ -170,16 +231,27 @@ bool doMergeSort(Phase1PipelineArg& p, Phase1BamHeap& heap) { void* phase1MergeSort(void* data) { Phase1PipelineArg& p = *(Phase1PipelineArg*)data; gUncompressDataBuf = p.uncompressData.dataBuf; // 主要用来比较bam的名字信息 + Phase1BamHeap posHeap; + Phase1BamHeap nameHeap; + if (nsgv::gSortArg.SORT_COORIDINATE) { + posHeap.Init(&p); + } else { + nameHeap.Init(&p); + } // merge - Phase1BamHeap heap; - heap.Init(&p); /* do the work */ while (true) { // previous dependency yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, p.MERGE_BUF_NUM); PROF_G_BEG(merge); - bool finish = doMergeSort(p, heap); + // bool finish = doMergeSort(p, heap); + bool finish = false; + if (nsgv::gSortArg.SORT_COORIDINATE) { + finish = doMergeSort(p, posHeap); + } else { + finish = doMergeSort(p, nameHeap); + } PROF_G_END(merge); if (finish) { diff --git a/src/sort/phase_1_sort.h b/src/sort/phase_1_sort.h index 2d5a183..86c5cfc 100644 --- a/src/sort/phase_1_sort.h +++ b/src/sort/phase_1_sort.h @@ -12,4 +12,9 @@ void phase1Sort(void* data); -void phase1MergeCompress(void* data); \ No newline at end of file +void phase1MergeCompress(void* data); + +void mtInThreadSort(void* data, long idx, int tid); +void singleThreadSort(void* data); + +void mergeLastData(void* p); \ No newline at end of file diff --git a/src/sort/phase_1_uncompress.cpp b/src/sort/phase_1_uncompress.cpp index 287ff25..6de5c5d 100644 --- a/src/sort/phase_1_uncompress.cpp +++ b/src/sort/phase_1_uncompress.cpp @@ -188,12 +188,13 @@ static void mtUncompressBlockBatch(void* data, long idx, int tid) { static void handleAdjacentThreadBlock(Phase1PipelineArg& p) { auto& uncompressData = p.uncompressData; auto& threadUncompressDataArr = p.threadUncompressWrap[p.uncompressOrder % p.UNCOMPRESS_BUF_NUM].threadUncompressDataArr; - size_t offset = 0; // 当前线程对应的全局数据的起始偏移量 - size_t bamOffset = 0; // 当前线程解析的bam在全局数据中的偏移量 + int64_t offset = 0; // 当前线程对应的全局数据的起始偏移量, 第一个可能是负数 + int64_t bamOffset = 0; // 当前线程解析的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; @@ -238,14 +239,25 @@ static void handleAdjacentThreadBlock(Phase1PipelineArg& p) { leftDataLen += additionDataLen; } // 重新解析 + blockBuf.readPos = leftDataLen; bamArr.Clear(); ParseAddAllBams(blockBuf.data, leftDataLen, blockBuf.curLen, bamArr, nullptr, &blockBuf.lastPos); spdlog::error("bam len mismatch {}: {}, {}, addition len: {}", tid, bamLen, claculatedBamLen, additionDataLen); } ParseAddBam(lastBamBuf.data, firstBam); } + //if (threadUncompressDataArr[tid].firstBam.Size() > 0) { + // threadUncompressDataArr[tid].memOffset -= threadUncompressDataArr[tid].firstBam.arr[0].wholeBamLen - blockBuf.readPos; + //} +#if 0 + if (hasLastData) + offset += lastBamBuf.curLen; + offset += threadUncompressDataArr[tid].blockBuf.curLen - threadUncompressDataArr[tid].blockBuf.readPos; + bamOffset += threadUncompressDataArr[tid].bamArr.Size() + threadUncompressDataArr[tid].firstBam.Size(); +#else offset += threadUncompressDataArr[tid].blockBuf.curLen; bamOffset += threadUncompressDataArr[tid].bamArr.Size() + threadUncompressDataArr[tid].firstBam.Size(); +#endif if (tid == p.numThread - 1) { // 最后一个线程 lastRoundBuf.Clear(); int lastBlockLeftDataLen = blockBuf.curLen - blockBuf.lastPos; @@ -257,7 +269,7 @@ static void handleAdjacentThreadBlock(Phase1PipelineArg& p) { } /* 将gz block进行解压，并进行线程内排序 */ -static void doPhase1Uncompress(Phase1PipelineArg& p, int finish = 0) { +static void doPhase1Uncompress(Phase1PipelineArg& p) { PROF_G_BEG(uncompress); kt_for(p.numThread, mtUncompressBlockBatch, &p, p.numThread); @@ -281,7 +293,7 @@ void* phase1Uncompress(void* data) { if (p.readFinish) { while (p.uncompressOrder < p.readOrder) { yarn::DEPENDENCY_NOT_TO_BE(p.uncompressSig, p.UNCOMPRESS_BUF_NUM); - doPhase1Uncompress(p, 1); + doPhase1Uncompress(p); yarn::UPDATE_SIG_ORDER(p.uncompressSig, p.uncompressOrder); } yarn::SIGNAL_FINISH(p.uncompressSig, p.uncompressFinish); @@ -343,19 +355,36 @@ static void mtMemCopy(void* data, long idx, int tid) { auto& uncompressData = p.uncompressData; // 所有线程共用一个，串行往这里添加解压后的block数据 // 拷贝bam未解析数据到全局的uncompressData里 +#if 0 + int firstBamDataLen = 0; + if (threadUncompressDataArr[tid].firstBam.Size() > 0) { + firstBamDataLen += threadUncompressDataArr[tid].lastBamBuf.curLen; + memcpy(uncompressData.dataBuf + uncompressData.usedBufSize + threadUncompressDataArr[tid].memOffset, + threadUncompressDataArr[tid].lastBamBuf.data, threadUncompressDataArr[tid].lastBamBuf.curLen); + } + memcpy( + uncompressData.dataBuf + uncompressData.usedBufSize + threadUncompressDataArr[tid].memOffset + threadUncompressDataArr[tid].lastBamBuf.curLen, + threadUncompressDataArr[tid].blockBuf.data + threadUncompressDataArr[tid].blockBuf.readPos, + threadUncompressDataArr[tid].blockBuf.curLen - threadUncompressDataArr[tid].blockBuf.readPos); +#else memcpy(uncompressData.dataBuf + uncompressData.usedBufSize + threadUncompressDataArr[tid].memOffset, threadUncompressDataArr[tid].blockBuf.data, threadUncompressDataArr[tid].blockBuf.curLen); +#endif // 拷贝解析的bam到全局数据里 size_t i = p.allBams.Size(); // 当前线程解析的bam在全局数据中的起始偏移量, 这时curIdx还没更新 uint64_t startOffset = uncompressData.usedBufSize + threadUncompressDataArr[tid].memOffset; // 当前线程解析的bam在全局数据中的起始偏移量 for (size_t j = 0; j < threadUncompressDataArr[tid].firstBam.Size(); ++i, ++j) { - p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].firstBam.arr[j]; - p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset].offset += startOffset; // 更新bam的偏移量 + size_t idx = i + threadUncompressDataArr[tid].bamOffset; + p.allBams.arr[idx] = threadUncompressDataArr[tid].firstBam.arr[j]; + p.allBams.arr[idx].offset += startOffset + threadUncompressDataArr[tid].blockBuf.readPos - p.allBams.arr[idx].wholeBamLen; // 更新bam的偏移量 + // CheckBam(uncompressData.dataBuf + p.allBams.arr[idx].offset, p.allBams.arr[idx].wholeBamLen); 前面线程数据可能还没拷贝好 } for (size_t j = 0; j < threadUncompressDataArr[tid].bamArr.Size(); ++i, ++j) { - p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset] = threadUncompressDataArr[tid].bamArr.arr[j]; - p.allBams.arr[i + threadUncompressDataArr[tid].bamOffset].offset += startOffset; // 更新bam的偏移量 + size_t idx = i + threadUncompressDataArr[tid].bamOffset; + p.allBams.arr[idx] = threadUncompressDataArr[tid].bamArr.arr[j]; + p.allBams.arr[idx].offset += startOffset; // 更新bam的偏移量 + // CheckBam(uncompressData.dataBuf + p.allBams.arr[idx].offset, p.allBams.arr[idx].wholeBamLen); } } @@ -366,13 +395,32 @@ static void doMemCopy(Phase1PipelineArg& p) { auto& lastThreadData = uncompressWrap.threadUncompressDataArr[p.numThread - 1]; size_t newDataLen = lastThreadData.memOffset + lastThreadData.blockBuf.curLen; + // 确保全局缓冲能放下这些解压的blocks + size_t dataSize = uncompressWrap.GetTotalBlockNum() * SINGLE_BLOCK_SIZE; + p.maxMemBytes = p.maxMemBytes < dataSize ? dataSize : p.maxMemBytes; + p.uncompressData.ReAllocate(p.maxMemBytes); + // 判断缓存是否已满 - size_t totalDataLen = p.uncompressData.usedBufSize + p.allBams.Size() * sizeof(OneBam) + newDataLen; + // size_t totalDataLen = p.uncompressData.usedBufSize + p.allBams.Size() * sizeof(OneBam) + newDataLen; + size_t totalDataLen = p.uncompressData.usedBufSize + newDataLen; + spdlog::info("total: {}, used: {}, new: {}", totalDataLen, p.uncompressData.usedBufSize, newDataLen); if (totalDataLen > p.maxMemBytes) { spdlog::info("block num: {}, all block num: {}, bam num: {}, all bam num: {}", uncompressWrap.GetTotalBlockNum(), p.blockNum, uncompressWrap.GetTotalBamNum(), p.bamNum); // 测试一下bam的offset对不对 +#if 0 + uint64_t allBamLen = 0; + for (int i = 0; i < p.allBams.Size(); ++i) { + CheckBam(p.uncompressData.dataBuf + p.allBams.arr[i].offset, p.allBams.arr[i].wholeBamLen); + + if (allBamLen != p.allBams.arr[i].offset) { + spdlog::info("i: {}, len-offset: {}-{}", i, allBamLen, p.allBams.arr[i].offset); + } + allBamLen += p.allBams.arr[i].wholeBamLen; + } +#endif + #if 0 auto &b = p.allBams.arr[p.allBams.Size() - 1]; // 最后一个bam OneBam o; @@ -411,6 +459,26 @@ static void doMemCopy(Phase1PipelineArg& p) { PROF_G_END(mem_copy); } +static void handleLastRoundData(Phase1PipelineArg& p) { + // if (true) { + if (p.uncompressData.usedBufSize * p.numThread < p.uncompressData.dataBufSize * 2) { + // 此时直接单线程排序就行了 + PROF_G_BEG(sort); + singleThreadSort(&p); + PROF_G_END(sort); + } else { + // 多线程排序，再归并 + PROF_G_BEG(sort); + kt_for(p.numThread, mtInThreadSort, &p, p.numThread); + PROF_G_END(sort); + + // 归并 + PROF_G_BEG(merge); + mergeLastData(&p); + PROF_G_END(merge); + } +} + /* phase1Uncompress step-3 拷贝线程 */ void* phase1MemCopy(void* data) { Phase1PipelineArg& p = *(Phase1PipelineArg*)data; @@ -426,7 +494,8 @@ void* phase1MemCopy(void* data) { } // 这里需要再检查一次缓冲区，有数据的话需要处理，只需要线程内排序，再线程间归并排序，不需要压缩写入中间文件了 - // todo + if (p.allBams.Size() > 0) + handleLastRoundData(p); break; } diff --git a/src/sort/phase_1_write.cpp b/src/sort/phase_1_write.cpp index 4bc6dd7..110c9ef 100644 --- a/src/sort/phase_1_write.cpp +++ b/src/sort/phase_1_write.cpp @@ -11,6 +11,7 @@ #include #include +#include #include #include @@ -23,15 +24,52 @@ #include "sort.h" #include "util/profiling.h" +static void checkCompress(uint8_t *addr, uint64_t len) { + size_t curReadPos = 0; + int blockLen = 0; + int maxBlockLen = 0; + DataBuffer buf; + buf.AllocMem(SINGLE_BLOCK_SIZE); + + while (curReadPos + BLOCK_HEADER_LENGTH <= len) { /* 确保能解析block长度 */ + blockLen = unpackInt16(&addr[curReadPos + 16]) + 1; + if (blockLen > maxBlockLen) { + maxBlockLen = blockLen; + } + if (curReadPos + blockLen <= len) { /* 完整的block数据在buf里 */ + size_t dlen = SINGLE_BLOCK_SIZE; // 65535 + uint32_t crc = le_to_u32(addr + curReadPos + blockLen - 8); + int ret = bgzfUncompress(buf.data, &dlen, (Bytef*)(addr + curReadPos) + BLOCK_HEADER_LENGTH, blockLen - BLOCK_HEADER_LENGTH, crc); + if (ret != 0) { + spdlog::error("block len: {}, uncompressed len: {}", blockLen, dlen); + exit(0); + } + + curReadPos += blockLen; + } else { + spdlog::error("not valid compressed block: {}, {}", curReadPos + blockLen, len); + break; /* 当前block数据不完整，一部分在还没读入的file数据里 */ + } + } + if (curReadPos != len) { + spdlog::error("addr: {}, len: {}", curReadPos, len); + exit(0); + } +} + static void doWrite(Phase1PipelineArg& p) { PROF_G_BEG(write_mid); - DataBuffer& compressBuf = p.compressBuf[p.compressOrder % p.COMPRESS_BUF_NUM]; + DataBuffer& compressBuf = p.compressBuf[p.writeOrder % p.COMPRESS_BUF_NUM]; + // checkCompress(compressBuf.data, compressBuf.curLen); fwrite(compressBuf.data, 1, compressBuf.curLen, p.midFilePtr); PROF_G_END(write_mid); } void* phase1Write(void* data) { Phase1PipelineArg& p = *(Phase1PipelineArg*)data; + // for test，写header + // fwrite(nsgv::gInHdr.compressed.data, 1, nsgv::gInHdr.compressed.curLen, p.midFilePtr); + /* do the work */ while (true) { // previous dependency @@ -50,6 +88,9 @@ void* phase1Write(void* data) { p.writeOrder += 1; } + // 写结尾，空block，bam文件需要 + //fwrite("\037\213\010\4\0\0\0\0\0\377\6\0\102\103\2\0\033\0\3\0\0\0\0\0\0\0\0\0", 1, 28, p.midFilePtr); + spdlog::info("End write order: {}", p.writeOrder); return nullptr; } \ No newline at end of file diff --git a/src/sort/phase_2.cpp b/src/sort/phase_2.cpp index f753149..bf6d89f 100644 --- a/src/sort/phase_2.cpp +++ b/src/sort/phase_2.cpp @@ -9,6 +9,31 @@ #include "phase_2.h" -void phase2Pipeline(Phase2PipelineArg& p) { +#include +#include +#include "common_data.h" +#include "phase_2_read.h" +#include "phase_2_merge.h" +#include "phase_2_write.h" +#include "util/profiling.h" + + +void phase2Pipeline(Phase2PipelineArg& p) { + PROF_G_BEG(phase2); + /* create threads */ + pthread_t tidArr[6]; // 2-stage pipeline + pthread_create(&tidArr[0], NULL, phase2ReadMidFile, &p); + pthread_create(&tidArr[1], NULL, phase2Uncompress, &p); + pthread_create(&tidArr[2], NULL, phase2CopyToMergeBuf, &p); + pthread_create(&tidArr[3], NULL, phase2Merge, &p); + pthread_create(&tidArr[4], NULL, phase2Compress, &p); + pthread_create(&tidArr[5], NULL, phase2Write, &p); + + for (int i = 0; i < 6; ++i) pthread_join(tidArr[i], NULL); + //for (int i = 0; i < 4; ++i) pthread_join(tidArr[i], NULL); + + spdlog::info("all bams num: {}", p.numBam); + + PROF_G_END(phase2); } \ No newline at end of file diff --git a/src/sort/phase_2.h b/src/sort/phase_2.h index 63fef8c..5ca6e53 100644 --- a/src/sort/phase_2.h +++ b/src/sort/phase_2.h @@ -13,32 +13,539 @@ */ #pragma once +#include #include +#include -struct Phase2File { - FILE* fp; - // 双buffer - // 当前读入的buffer指针df - // ReadBuffer - -}; +#include "phase_1.h" +#include "sam_io.h" +#include "sort.h" +using std::string; // 循环缓冲区 struct CircularBuffer { - + uint8_t* data = nullptr; + size_t readIdx = 0; // 可以读取的开始位置 + size_t writeIdx = 0; // 可以写入的开始位置 + size_t valueSize = 0; // 有效字节 + size_t bufSize = 0; // 缓冲区空间 + + CircularBuffer() {} + CircularBuffer(size_t initSize) { + bufSize = initSize; + data = (uint8_t*)malloc(bufSize); + } + ~CircularBuffer() { + if (data) { + free(data); + data = nullptr; + } + readIdx = 0; + writeIdx = 0; + valueSize = 0; + bufSize = 0; + } + void AllocMem(size_t memSize) { ReAllocMem(memSize); } + + void ReAllocMem(size_t memSize) { + if (memSize > bufSize) { + bufSize = memSize; + data = (uint8_t*)realloc(data, bufSize); + } + } + + size_t Write(uint8_t* src, size_t len) { + if (len == 0) return 0; + size_t freeSpace = bufSize - valueSize; + size_t toWrite = MIN(len, freeSpace); + if (toWrite == 0) return 0; + + size_t firstPart = MIN(toWrite, bufSize - writeIdx); + // 第一段：从 write_index_ 到末尾 + memcpy(data + writeIdx, src, firstPart); + // 第二段：如果超过了末尾，则从头部继续 + if (toWrite > firstPart) { + memcpy(data, src + firstPart, toWrite - firstPart); + } + writeIdx = (writeIdx + toWrite) % bufSize; + valueSize += toWrite; + return toWrite; + } + + size_t Read(uint8_t* out, size_t len) { + if (len == 0) return 0; + size_t toRead = MIN(len, valueSize); + if (toRead == 0) return 0; + + size_t firstPart = MIN(toRead, bufSize - readIdx); + memcpy(out, data + readIdx, firstPart); + if (toRead > firstPart) { + std::memcpy(out + firstPart, data, toRead - firstPart); + } + readIdx = (readIdx + toRead) % bufSize; + valueSize -= toRead; + return toRead; + } + + // 返回第一个连续空间的大小 + size_t FirstPartWriteSize() { + size_t freeSpace = bufSize - valueSize; + size_t firstPart = MIN(freeSpace, bufSize - writeIdx); + return firstPart; + } + + // 如果空间不连续，那么返回第二个连续空间的内存大小 + size_t SecondPartWriteSize() { + size_t freeSpace = bufSize - valueSize; + size_t firstPart = MIN(freeSpace, bufSize - writeIdx); + if (firstPart == freeSpace) + return 0; + return freeSpace - firstPart; + } + + // 跳过不能完整保存一个bam的空间 + void SkipWrite(size_t skipBytes) { + valueSize += skipBytes; + writeIdx = (writeIdx + skipBytes) % bufSize; + } + + void SkipRead(size_t from, size_t skipBytes) { + if (valueSize <= skipBytes) { // 直接清空 + readIdx = 0; // 可以读取的开始位置 + writeIdx = 0; // 可以写入的开始位置 + valueSize = 0; // 有效字节 + } else { + if (from == readIdx) { + readIdx = (readIdx + skipBytes) % bufSize; + valueSize -= skipBytes; + } else { + valueSize -= (bufSize - readIdx + skipBytes); + readIdx = skipBytes; + } + } + } + + // 退回一个bam + void RevertRead(size_t from, size_t revertBytes) { + readIdx = from - revertBytes; + valueSize += revertBytes; + } + + // 读取一个bam + size_t ReadBam(uint8_t* out, size_t start, size_t len) { + memcpy(out, data + start, len); + valueSize -= len; + if (readIdx != start) { + valueSize -= bufSize - readIdx; + } + readIdx = start + len; + return len; + } + + // 是否为空 + bool Empty() const { return valueSize == 0; } + // 是否已满 + bool Full() const { return valueSize == bufSize; } + // 当前数据量（字节） + size_t Size() const { return valueSize; } + // 剩余可写空间 + size_t Free() const { return bufSize - valueSize; } + // 总容量 + size_t Capacity() const { return bufSize; } + + void Clear() { + readIdx = 0; // 可以读取的开始位置 + writeIdx = 0; // 可以写入的开始位置 + valueSize = 0; // 有效字节 + } }; - // 循环数组 +template struct CircularArray { + std::vector arr; + size_t readIdx = 0; // 可以读取的开始位置 + size_t writeIdx = 0; // 可以写入的开始位置 + size_t valueSize = 0; // 有效字节 + size_t bufSize = 0; // 缓冲区空间 + CircularArray() {} + CircularArray(size_t initSize) { + bufSize = initSize; + arr.resize(initSize); + } + void AllocArr(size_t arrSize) { ReAllocArr(arrSize); } + + void ReAllocArr(size_t arrSize) { + if (arrSize > bufSize) { + bufSize = arrSize; + arr.resize(arrSize); + } + } + + inline T& operator[](size_t idx) { + return arr[(idx + readIdx) % bufSize]; + } + + // 推入元素（自动扩容） + void Push(const T& value) { + if (Full()) { + ReAllocArr(bufSize * 1.5); + } + arr[writeIdx] = value; + writeIdx = (writeIdx + 1) % bufSize; + ++valueSize; + } + + void Push(T&& value) { + if (Full()) { + ReAllocArr(bufSize * 1.5); + } + arr[writeIdx] = std::move(value); + writeIdx = (writeIdx + 1) % bufSize; + ++valueSize; + } + + // 弹出头部元素（如果为空则返回 nullopt） + void Pop() { + if (Empty()) + return; + --valueSize; + readIdx = (readIdx + 1) % bufSize; + } + + void Revert() { + ++valueSize; + readIdx = (readIdx - 1) % bufSize; + } + + // 查看头部/尾部（不弹出） + T* Front() { + if (Empty()) + return nullptr; + return &arr[readIdx]; + } + + T* Back() { + if (Empty()) + return nullptr; + return &arr[(writeIdx - 1) % bufSize]; + } + + // 是否为空 + bool Empty() const { return valueSize == 0; } + // 是否已满 + bool Full() const { return valueSize == bufSize; } + // 当前数据量（字节） + size_t Size() const { return valueSize; } + // 剩余可写空间 + size_t Free() const { return bufSize - valueSize; } + // 总容量 + size_t Capacity() const { return bufSize; } + + + void Clear() { + readIdx = 0; // 可以读取的开始位置 + writeIdx = 0; // 可以写入的开始位置 + valueSize = 0; // 有效字节 + } +}; + +// 合并过程需要的缓冲区 +struct Phase2MergeBuffer { + CircularBuffer data; + CircularArray bams; + bool initialized = false; + size_t numProcessed = 0; + + OneBam* Front() { return bams.Front(); } + void Pop() { + OneBam* bp = Front(); + data.SkipRead(bp->offset, bp->wholeBamLen); + bams.Pop(); + } + + void Revert() { + bams.Revert(); + OneBam* b = bams.Front(); + data.RevertRead(b->offset, b->wholeBamLen); + } + + size_t Size() { return bams.Size(); } + + bool hasSpace(size_t bamLen) { return data.FirstPartWriteSize() >= bamLen; } + + // 从一个解压后的block缓冲区拷贝多个bam到循环缓冲区 + // 返回实际拷贝的bam数量 + size_t CopyBams(DataBuffer &blockBuf, BamArr &arr, int start) { + if (start >= arr.Size()) return 0; + int origStart = start; + size_t numCopied = 0; + OneBam* b1 = &arr.Get(start); + OneBam* b2 = arr.Back(); + + size_t firstPartSize = data.FirstPartWriteSize(); + size_t secondPartSize = data.SecondPartWriteSize(); + size_t needSize = b2->offset - b1->offset + b2->wholeBamLen; + int stop = arr.Size(); + if (needSize <= firstPartSize) { // 在第一个连续空间里就能放下 + // 每个bam的offset需要加上diff，以对应新的buf + int64_t diff = (int64_t)data.writeIdx - b1->offset; + data.Write(blockBuf.data + b1->offset, needSize); + for (int i = start; i < stop; ++i) { + bams.Push(arr.Get(i)); + bams.Back()->offset += diff; + bams.Back()->addr = data.data; + } + } else { + stop = start; + size_t firstNeedSize = 0; + while (arr.Get(stop).wholeBamLen + firstNeedSize < firstPartSize) { + firstNeedSize += arr.Get(stop).wholeBamLen; + stop += 1; + } + int64_t diff = (int64_t)data.writeIdx - b1->offset; + data.Write(blockBuf.data + b1->offset, firstNeedSize); + for (int i = start; i < stop; ++i) { + bams.Push(arr.Get(i)); + bams.Back()->offset += diff; + bams.Back()->addr = data.data; + } + // 跳过first part不能放下完整bam的部分 + data.SkipWrite(firstPartSize - firstNeedSize); + b1 = &arr.Get(stop); + start = stop; + // 拷贝第二段 + size_t secondNeedSize = needSize - firstNeedSize; + if (secondNeedSize <= secondPartSize) { + stop = arr.Size(); + } else { + secondNeedSize = 0; + while (arr.Get(stop).wholeBamLen + secondNeedSize < secondPartSize) { + secondNeedSize += arr.Get(stop).wholeBamLen; + stop += 1; + } + } + diff = (int64_t)data.writeIdx - b1->offset; + data.Write(blockBuf.data + b1->offset, secondNeedSize); + for (int i = start; i < stop; ++i) { + bams.Push(arr.Get(i)); + bams.Back()->offset += diff; + bams.Back()->addr = data.data; + } + } + numCopied = stop - origStart; + + return numCopied; + } + + void InitSize(size_t _bufsize, size_t _bamsize) { + data.AllocMem(_bufsize); + bams.AllocArr(_bamsize); + initialized = true; + } +}; + +struct BamDataBuffer { + uint64_t offset = 0; + int bamOffset = 0; + DataBuffer block; + BamArr bams; + void Clear() { + offset = 0; + block.Clear(); + bams.Clear(); + } +}; + +// 读取和解压需要的buffer +struct Phase2ReadBuffer : ReadBuffer { + FastVector uncompressBlocks; + int bamNum = 0; +}; + +// 解析 +struct UncompressBuffer { + DataBuffer blockBuf; // 解压的block放在这里 + BamArr bamArr; // 解析后的bam数据放在这里 + int startIdx = 0; + + size_t Size() { return bamArr.Size() - startIdx; } + + OneBam *Front() { + if (Size() == 0) + return nullptr; + return &bamArr.Get(startIdx); + } + + void Clear() { + startIdx = 0; + blockBuf.Clear(); + bamArr.Clear(); + } +}; + +struct Phase2File { + // static constexpr int MID_FILE_BUF_SIZE = 8 * SINGLE_BLOCK_SIZE; + static constexpr int READ_BUF_NUM = 2; // 读入的buf数量 + static constexpr int UNCOMPRESSS_BUF_NUM = 2; // 解析的buf数量 + static constexpr int COPY_BUF_NUM = 1; // 拷贝到merge缓冲区数量 + + // 参数 + size_t readBufSize = 0; + string fileName; + + FILE* fp = nullptr; + Phase2ReadBuffer readData[READ_BUF_NUM]; + UncompressBuffer uncompressBuf[UNCOMPRESSS_BUF_NUM]; + DataBuffer halfBlock; // 剩余不完整的压缩的block数据 + + // 读 + volatile int readyReadBufNum = 0; // 有几个buffer已经填满可用了, 类似lock_t的作用 + uint64_t readOrder = 0; // 当前读取的buf索引 + + // 解压 + volatile int readyUncompressBufNum = 0; + uint64_t uncompressOrder = 0; // 写入的buf索引，类似readOrder + bool needUncompress = false; + + // 拷贝到归并buf + Phase2MergeBuffer mergeData; // 归并时候用到的数据 + uint64_t copyOrder = 0; + volatile int readyMergeBufNum = 0; + + volatile int finish = 0; // 读入结束 + + void Init() { + for (int i = 0; i < READ_BUF_NUM; ++i) { + readData[i].Resize(readBufSize); + } + halfBlock.AllocMem(SINGLE_BLOCK_SIZE); + // 打开中间文件 + // spdlog::info("file: {}", fileName); + fp = fopen(fileName.c_str(), "rb"); + } + + void Init(const string& _fileName, size_t _readBufSize) { + readBufSize = _readBufSize; + fileName = _fileName; + Init(); + } + + // 缓冲区大小应该根据中间文件的数量进行调整 + + Phase2File() {} + Phase2File(const string &_fileName, size_t _readBufSize) { + readBufSize = _readBufSize; + fileName = _fileName; + } + + ~Phase2File() { + if (fp) { + fclose(fp); + } + } +}; + +// 用于合并压缩的数据结构 +struct Phase2MergeData { + vector blockDataArr; // 待压缩的数据 + vector compressDataArr; // 压缩后的数据 + + void Resize(int blockNum) { + blockDataArr.resize(blockNum); + compressDataArr.resize(blockNum); + } }; /* 第二阶段的多线程流水线参数 */ struct Phase2PipelineArg { + // common parameters + static const int MERGE_BUF_NUM = 2; // 压缩的buf数量 + static const int COMPRESS_BUF_NUM = 2; // 压缩的buf数量 + + int compressBlocksThreshold = 0; // 进行归并排序的，单次block数量阈值，参数 + int numThread = 0; // 线程数 + int compressLevel = 3; // 压缩级别，参数 + uint64_t numBam = 0; + + // 从phase1转移来的数据 + UncompressBlockBuffer& uncompressData; // 所有线程共用一个，串行往这里添加解压后的block数据 + BamArr& allBams; // 所有线程共用一个，串行往这里添加解析后的bam数据 + + // 读入中间文件 + vector midFiles; + yarn::lock_t* readSig; + uint64_t readOrder = 0; + volatile int readFinish = 0; + + // 解压读入的blocks + yarn::lock_t* uncompressSig; + uint64_t uncompressOrder = 0; + volatile int uncompressFinish = 0; + int curBlockNum = 0; + + // 拷贝到merge缓冲区 + yarn::lock_t* copyMerge; + uint64_t copyMergeOrder = 0; + int uncompressReadyNum = 0; + volatile int copyMergeFinish = 0; + + // 将解压的数据拷贝进merge的缓冲区，归并后放入，以便后续压缩用到 + MergeCompressData mergeData[MERGE_BUF_NUM]; + yarn::lock_t* mergeSig; + yarn::lock_t* copyCompressSig; + uint64_t mergeOrder = 0; // 前半段 + uint64_t copyCompressOrder = 0; // 后半段，需要压缩 + volatile int mergeFinish = 0; + volatile int copyCompressFinish = 0; + + // 压缩 + DataBuffer compressBuf[COMPRESS_BUF_NUM]; // 压缩用到的buf + yarn::lock_t* compressSig; + uint64_t compressOrder = 0; + volatile int compressFinish = 0; + + // 写入文件 + FILE* outFilePtr = nullptr; // 文件指针 + uint64_t writeOrder = 0; + + void SetThreadNum(int nThread) { + numThread = nThread; + // const int kThres = 1024; + const int kThres = 2048; + int threadBlocksThres = numThread * 32; + compressBlocksThreshold = threadBlocksThres < kThres ? threadBlocksThres : kThres; // 大概64M的解压缩后的数据 + for (int i = 0; i < COMPRESS_BUF_NUM; ++i) { + mergeData[i].Resize(compressBlocksThreshold); + } + } + + Phase2PipelineArg(UncompressBlockBuffer& lastBuf, BamArr& lastBams, string &midFilePrefix, int midFileNum, int threadNum) : uncompressData(lastBuf), allBams(lastBams){ + SetThreadNum(threadNum); + + int allFileBufSize = 64L * 1024 * 1024; + // int bufBlocks = (numThread * 4 + midFileNum) / (midFileNum + 1); + // size_t kFileBufSize = bufBlocks * SINGLE_BLOCK_SIZE; + size_t kFileBufSize = allFileBufSize / (midFileNum + 1); + + spdlog::info("mid file buf size: {}", kFileBufSize); + + midFiles.resize(midFileNum); + for (int i = 0; i < midFileNum; ++i) { + midFiles[i].Init(midFilePrefix + std::to_string(i), kFileBufSize); + } + + readSig = yarn::NEW_LOCK(0); + uncompressSig = yarn::NEW_LOCK(0); + copyMerge = yarn::NEW_LOCK(0); + mergeSig = yarn::NEW_LOCK(0); + copyCompressSig = yarn::NEW_LOCK(0); + compressSig = yarn::NEW_LOCK(0); + } }; void phase2Pipeline(Phase2PipelineArg &p); \ No newline at end of file diff --git a/src/sort/phase_2_merge.cpp b/src/sort/phase_2_merge.cpp index e69de29..9af2c66 100644 --- a/src/sort/phase_2_merge.cpp +++ b/src/sort/phase_2_merge.cpp @@ -0,0 +1,309 @@ +/* + Description: 归并中间文件 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/05 +*/ + +#include "phase_2_merge.h" + +#include +#include + +#include "common_data.h" +#include "phase_2.h" +#include "util/profiling.h" + +/* bam 排序堆 */ +struct Phase2BamArrIdIdx { + int idx = 0; // 如果是第一阶段转移过来的buffer，用这个，此时file为nullptr + uint8_t* addr = nullptr; // bam的原始地址，就是减去offset + Phase2File* file = nullptr; + const OneBam* bam = nullptr; +}; + +struct Phase2BamPosGreaterThan { + bool operator()(const Phase2BamArrIdIdx& a, const Phase2BamArrIdIdx& b) const { return a.bam->pos > b.bam->pos; } +}; + +struct Phase2BamNameGreaterThan { + bool operator()(const Phase2BamArrIdIdx& a, const Phase2BamArrIdIdx& b) const { + int cmp = strncmp((char*)(a.addr + a.bam->offset + OneBam::QnameOffset), + (char*)(b.addr + b.bam->offset + OneBam::QnameOffset), std::min(a.bam->qnameLen, b.bam->qnameLen)); + if (cmp == 0) + return a.bam->qnameLen > b.bam->qnameLen; + return cmp > 0; + } +}; + +/* 用来排序 bam*/ +template +struct Phase2BamHeap { + Phase2PipelineArg* p; + priority_queue, GreaterThan> minHeap; + size_t popNum = 0; + size_t curIdx = 0; + + int Init(Phase2PipelineArg* _p) { + p = _p; + if (p == nullptr) { + return -1; + } + // for (int i = 0; i < 4; ++i) { + for (int i = 0; i < p->midFiles.size(); ++i) { + auto& f = p->midFiles[i]; + f.mergeData.Front()->addr = f.mergeData.data.data; + if (f.mergeData.Front()->addr == nullptr) { + spdlog::error("null addr"); + } + minHeap.push({0, f.mergeData.data.data, &f, f.mergeData.Front()}); + f.mergeData.Pop(); + } + if (p->allBams.Size() > curIdx) { + p->allBams.arr[curIdx].addr = p->uncompressData.dataBuf; + if (p->allBams.arr[curIdx].addr == nullptr) { + spdlog::error("null addr"); + } + minHeap.push({0, p->uncompressData.dataBuf, nullptr, &p->allBams.arr[curIdx]}); + } + return 0; + } + + void ReInitFile(Phase2File *file) { + auto &f = *file; + if (f.mergeData.Front() != nullptr) { + f.mergeData.Front()->addr = f.mergeData.data.data; + if (f.mergeData.Front()->addr == nullptr) { + spdlog::error("null addr"); + } + minHeap.push({0, f.mergeData.data.data, &f, f.mergeData.Front()}); + f.mergeData.Pop(); + } + } + + void Revert() { + while (!minHeap.empty()) { + auto minVal = minHeap.top(); + if (minVal.file != nullptr) { + auto& mergeData = minVal.file->mergeData; + mergeData.Revert(); + } else { + --curIdx; + } + minHeap.pop(); + } + } + + const OneBam* Pop() { + const OneBam* ret = nullptr; + if (!minHeap.empty()) { + auto minVal = minHeap.top(); + minHeap.pop(); + ++popNum; + ret = minVal.bam; + if (minVal.file != nullptr) { + auto& mergeData = minVal.file->mergeData; + if (mergeData.Size() > 0) { + mergeData.Front()->addr = minVal.addr; + if (mergeData.Front()->addr == nullptr) { + spdlog::error("null addr"); + } + minHeap.push({0, minVal.addr, minVal.file, mergeData.Front()}); + mergeData.Pop(); + } + } else { + if (minVal.idx + 1 < p->allBams.Size()) { + curIdx = minVal.idx + 1; + p->allBams.arr[minVal.idx + 1].addr = minVal.addr; + minHeap.push({minVal.idx + 1, minVal.addr, nullptr, &p->allBams.arr[minVal.idx + 1]}); + } + } + } + return ret; + } + + const OneBam* Pop(Phase2File ** emptyFile) { + const OneBam* ret = nullptr; + if (!minHeap.empty()) { + auto minVal = minHeap.top(); + minHeap.pop(); + ++popNum; + ret = minVal.bam; + if (minVal.file != nullptr) { + auto& mergeData = minVal.file->mergeData; + if (mergeData.Size() > 0) { + mergeData.Front()->addr = minVal.addr; + if (mergeData.Front()->addr == nullptr) { + spdlog::error("null addr"); + } + minHeap.push({0, minVal.addr, minVal.file, mergeData.Front()}); + mergeData.Pop(); + } else { + *emptyFile = minVal.file; + } + } else { + if (minVal.idx + 1 < p->allBams.Size()) { + curIdx = minVal.idx + 1; + p->allBams.arr[minVal.idx + 1].addr = minVal.addr; + if (p->allBams.arr[minVal.idx + 1].addr == nullptr) { + spdlog::error("null addr"); + } + minHeap.push({minVal.idx + 1, minVal.addr, nullptr, &p->allBams.arr[minVal.idx + 1]}); + } + } + } + return ret; + } + + const OneBam* Top() { + const OneBam* ret = nullptr; + if (!minHeap.empty()) { + ret = minHeap.top().bam; + } + return ret; + } +}; + +static void mtCopyBams(void* data, long idx, int tid) { + Phase2PipelineArg& p = *(Phase2PipelineArg*)data; + MergeCompressData& mergeData = p.mergeData[p.mergeOrder % p.MERGE_BUF_NUM]; + + auto& bams = mergeData.blockDataArr[idx].bamPtrArr; + auto& blockData = mergeData.blockDataArr[idx].blockBuf; + + // spdlog::info("bam size: {}", bams.Size()); + for (int i = 0; i < bams.Size(); ++i) { + const OneBam* bp = bams.arr[i]; + blockData.MemCopy(bp->addr + bp->offset, bp->wholeBamLen); + // check bam + // CheckBam(p.uncompressData.dataBuf + bp->offset, bp->wholeBamLen); + } +} + +template +static bool doPhase2Merge(Phase2PipelineArg& p, Phase2BamHeap& heap, bool *mergeFull = nullptr, Phase2File **emptyFilePtr = nullptr, bool copyFinish = false) { + bool finish = false; + + auto& mergeData = p.mergeData[p.mergeOrder % p.MERGE_BUF_NUM]; + const OneBam* bam = nullptr; + + size_t bamBytes = 0; + int singleBlockBytes = 0xff00; + Phase2File* emptyFile = nullptr; + + if (emptyFilePtr != nullptr && *emptyFilePtr != nullptr) { + heap.Init(&p); + } + + mergeData.blockDataArr[mergeData.curIdx].Clear(); + while ((bam = heap.Top()) != nullptr) { + if (bam->addr == nullptr) { + spdlog::info("null addr"); + } + if (bamBytes + bam->wholeBamLen > singleBlockBytes) { + mergeData.curIdx++; + if (mergeData.curIdx >= p.compressBlocksThreshold) { + if (mergeFull != nullptr) + *mergeFull = true; + break; + } + mergeData.blockDataArr[mergeData.curIdx].Clear(); // 清理，为添加bam数据做准备 + bamBytes = 0; + } +#if 0 + mergeData.blockDataArr[mergeData.curIdx].blockBuf.MemCopy(p.uncompressData.dataBuf + bam->offset, bam->wholeBamLen); +#else + mergeData.blockDataArr[mergeData.curIdx].bamPtrArr.Add(bam); +#endif + bamBytes += bam->wholeBamLen; // for test + if (copyFinish) { + heap.Pop(); + } else { + heap.Pop(&emptyFile); + if (emptyFile != nullptr) + break; + } + } + // 并行拷贝bam数据 + kt_for(p.numThread, mtCopyBams, &p, mergeData.curIdx); + + if (bam == nullptr) { + // 都处理完 + finish = true; + } + + if (emptyFile != nullptr) { + heap.Revert(); + } + + if (emptyFilePtr != nullptr) { + *emptyFilePtr = emptyFile; + } + + return finish; +} + +void* phase2Merge(void* data) { + + Phase2PipelineArg& p = *(Phase2PipelineArg*)data; + Phase2BamHeap posHeap; + Phase2BamHeap nameHeap; + bool firstInit = true; + Phase2File* emptyFile = nullptr; +#if 1 + // merge + /* do the work */ + while (true) { + // previous dependency + yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, p.MERGE_BUF_NUM); + yarn::DEPENDENCY_NOT_TO_BE(p.copyMerge, 0); + + // bool finish = doMergeSort(p, heap); + bool finish = false; + bool mergeFull = false; + emptyFile = nullptr; + + if (firstInit) { + if (nsgv::gSortArg.SORT_COORIDINATE) { + posHeap.Init(&p); + } else { + nameHeap.Init(&p); + } + firstInit = false; + } + if (p.copyMergeFinish) { + while (!finish) { + PROF_G_BEG(phase2_merge); + yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, p.MERGE_BUF_NUM); + if (nsgv::gSortArg.SORT_COORIDINATE) { + finish = doPhase2Merge(p, posHeap, &mergeFull, &emptyFile, true); + } else { + finish = doPhase2Merge(p, nameHeap, &mergeFull, &emptyFile, true); + } + if (mergeFull) + yarn::UPDATE_SIG_ORDER(p.mergeSig, p.mergeOrder); + PROF_G_END(phase2_merge); + } + yarn::SIGNAL_FINISH(p.mergeSig, p.mergeFinish); + break; + } + PROF_G_BEG(phase2_merge); + if (nsgv::gSortArg.SORT_COORIDINATE) { + doPhase2Merge(p, posHeap, &mergeFull, &emptyFile); + } else { + doPhase2Merge(p, nameHeap, &mergeFull, &emptyFile); + } + PROF_G_END(phase2_merge); + + // update self status + if (emptyFile != nullptr) // 只有某个文件的buf消耗完了，才读入 + yarn::CONSUME_SIGNAL(p.copyMerge); + if (mergeFull) + yarn::UPDATE_SIG_ORDER(p.mergeSig, p.mergeOrder); + } +#endif + spdlog::info("End phase2 merge sort order: {}", p.mergeOrder); + return nullptr; +} \ No newline at end of file diff --git a/src/sort/phase_2_merge.h b/src/sort/phase_2_merge.h index e69de29..128c82f 100644 --- a/src/sort/phase_2_merge.h +++ b/src/sort/phase_2_merge.h @@ -0,0 +1,12 @@ +/* + Description: 归并中间文件 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/05 +*/ + +#pragma once + +void* phase2Merge(void* data); \ No newline at end of file diff --git a/src/sort/phase_2_read.cpp b/src/sort/phase_2_read.cpp index e69de29..505e8f4 100644 --- a/src/sort/phase_2_read.cpp +++ b/src/sort/phase_2_read.cpp @@ -0,0 +1,410 @@ +/* + Description: 中间文件读入和解析等 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/05 +*/ + +#include "phase_2_read.h" + +#include +#include +#include +#include +#include + +#include "common_data.h" +#include "const_val.h" +#include "phase_1_read.h" +#include "phase_2.h" +#include "sam_io.h" +#include "sort.h" +#include "util/profiling.h" +#include "util/yarn.h" + +//////////////////////////////////////////////// 读入中间文件 +extern size_t ReadBamFile(ReadBuffer& readData, DataBuffer& halfBlock, FILE* fpr); + +bool doReadMidFiles(Phase2PipelineArg& p) { + PROF_G_BEG(phase2_read); + bool finish = false; + int finishNum = 0; + for (int i = 0; i < p.midFiles.size(); ++i) { + auto& f = p.midFiles[i]; + if (f.readyReadBufNum < Phase2File::READ_BUF_NUM && !f.finish) { // 还有空余缓冲区 + auto& readData = f.readData[f.readOrder % Phase2File::READ_BUF_NUM]; + size_t readState = ReadBamFile(readData, f.halfBlock, f.fp); + if (readState == 0) { + f.finish = 1; + } else { + f.readOrder += 1; + f.readyReadBufNum += 1; + readData.uncompressBlocks.ClearResize(readData.startAddrArr.size()); + } + //spdlog::info("read size: {}", readState); + } + finishNum += f.finish; + //spdlog::info("fid: {}, readyNum: {}", i, f.readyReadBufNum); + } + //spdlog::info("phase2 read order: {}, finishNum: {}, files: {}", p.readOrder, finishNum, p.midFiles.size()); + finish = finishNum == p.midFiles.size(); + PROF_G_END(phase2_read); + return finish; +} + +void* phase2ReadMidFile(void* data) { + Phase2PipelineArg& p = *(Phase2PipelineArg*)data; + + /* do the work */ + while (true) { + // previous dependency + yarn::DEPENDENCY_NOT_TO_BE(p.readSig, Phase2File::READ_BUF_NUM); + PROF_G_BEG(merge); + bool finish = doReadMidFiles(p); + PROF_G_END(merge); + + if (finish) { + yarn::SIGNAL_FINISH(p.readSig, p.readFinish); + break; + } + // update self status + yarn::UPDATE_SIG_ORDER(p.readSig, p.readOrder); + } + + spdlog::info("End read mid file order: {}", p.readOrder); + return nullptr; +} + +//////////////////////////////////////////////// 解析中间文件 + +extern size_t ParseAddAllBams(uint8_t* dataAddr, size_t startOffset, size_t endOffset, BamArr& bamArr, size_t* nextBamStartPtr = nullptr, + size_t* lastPosPtr = nullptr); + +static void mtUncompressBlockBatch(void* data, long idx, int tid) { + //PROF_T_BEG(mem_copy); + + Phase2PipelineArg& p = *(Phase2PipelineArg*)data; + + tid = idx; // 静态分配任务，此时用idx代替tid + int startIdx = START_IDX(tid, p.numThread, p.curBlockNum); + int stopIdx = STOP_IDX(tid, p.numThread, p.curBlockNum); + + int blockNum = 0; + int startFileIdx = -1; + int stopFileIdx = 0; + // 确定范围 + for (int i = 0; i < p.midFiles.size(); ++i) { + auto& f = p.midFiles[i]; + if (f.needUncompress) { // 还有空余缓冲区 + blockNum += f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM].startAddrArr.size(); + } + if (startFileIdx == -1 && startIdx < blockNum) { startFileIdx = i; } + if (blockNum >= stopIdx) { + stopFileIdx = i; // 包含 + break; + } + } + blockNum = 0; + for (int i = 0; i < startFileIdx; ++i) { + auto& f = p.midFiles[i]; + if (f.needUncompress) { // 还有空余缓冲区 + auto& readData = f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM]; + blockNum += readData.startAddrArr.size(); + } + } + for (int i = startFileIdx; i <= stopFileIdx; ++i) { + auto& f = p.midFiles[i]; + int start = 0; + int stop = 0; + if (f.needUncompress) { // 还有空余缓冲区 + auto& readData = f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM]; + start = MAX(0, startIdx - blockNum); + stop = MIN(readData.startAddrArr.size(), stopIdx - blockNum); + + // 解压并解析block + for (int j = start; j < stop; ++j) { + auto& blockData = readData.uncompressBlocks.arr[j]; + blockData.Clear(); + auto& blockBuf = blockData.block; + blockBuf.ReAllocMem(SINGLE_BLOCK_SIZE); + + uint8_t* block = readData.startAddrArr[j]; + size_t dlen = SINGLE_BLOCK_SIZE; // 65535 + int block_length = unpackInt16(&block[16]) + 1; + uint32_t crc = le_to_u32(block + block_length - 8); + //if (block_length > 30000) + // spdlog::error("block len: {}", block_length); + int ret = bgzfUncompress(blockBuf.data + blockBuf.curLen, &dlen, (Bytef*)block + BLOCK_HEADER_LENGTH, + block_length - BLOCK_HEADER_LENGTH, crc); + if (ret != 0) { + spdlog::error("uncompress error, block id: {}, len: {}, ret: {}", idx, block_length, ret); + exit(0); + } + blockBuf.curLen += dlen; + + // 解析 + ParseAddAllBams(blockBuf.data, 0, blockBuf.curLen, blockData.bams); + } + + //spdlog::info("tid: {}, blockNum: {}, fid: {}, blocks: {}-{}, block range: {}-{}, {}", tid, blockNum, i, start, stop, startIdx, stopIdx, + // stopIdx - blockNum); + + blockNum += readData.startAddrArr.size(); + } + } +} + +// 计算对应的每个block在uncompress缓冲区的偏移量，并准备好那么多的空间 +static void calculateOffset(Phase2PipelineArg& p) { + for (int i = 0; i < p.midFiles.size(); ++i) { + uint64_t offset = 0; + int bamOffset = 0; + auto& f = p.midFiles[i]; + auto& readData = f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM]; + for (int j = 0; j < readData.uncompressBlocks.Size(); ++j) { + auto& block = readData.uncompressBlocks.arr[j]; + block.offset = offset; + block.bamOffset = bamOffset; + offset += block.block.curLen; + bamOffset += block.bams.Size(); + } + readData.bamNum = bamOffset; + + if (f.needUncompress) { + auto& uncompressData = f.uncompressBuf[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM]; + uncompressData.Clear(); + uncompressData.bamArr.ReSize(bamOffset); + uncompressData.blockBuf.ReSize(offset); + } + } +} + +static void mtCopyToUncompressBuf(void* data, long idx, int tid) { + Phase2PipelineArg& p = *(Phase2PipelineArg*)data; + + tid = idx; // 静态分配任务，此时用idx代替tid + int startIdx = START_IDX(tid, p.numThread, p.curBlockNum); + int stopIdx = STOP_IDX(tid, p.numThread, p.curBlockNum); + + int blockNum = 0; + int startFileIdx = -1; + int stopFileIdx = 0; + // 确定范围 + for (int i = 0; i < p.midFiles.size(); ++i) { + auto& f = p.midFiles[i]; + if (f.needUncompress) { // 还有空余缓冲区 + blockNum += f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM].startAddrArr.size(); + } + if (startFileIdx == -1 && startIdx < blockNum) { + startFileIdx = i; + } + if (blockNum >= stopIdx) { + stopFileIdx = i; // 包含 + break; + } + } + blockNum = 0; + for (int i = 0; i < startFileIdx; ++i) { + auto& f = p.midFiles[i]; + if (f.needUncompress) { // 还有空余缓冲区 + auto& readData = f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM]; + blockNum += readData.startAddrArr.size(); + } + } + + for (int i = startFileIdx; i <= stopFileIdx; ++i) { + auto& f = p.midFiles[i]; + int start = 0; + int stop = 0; + if (f.needUncompress) { // 还有空余缓冲区 + auto& readData = f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM]; + auto& uncompressData = f.uncompressBuf[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM]; + + start = MAX(0, startIdx - blockNum); + stop = MIN(readData.startAddrArr.size(), stopIdx - blockNum); + + // 拷贝 + for (int j = start; j < stop; ++j) { + auto& blockData = readData.uncompressBlocks.arr[j]; + auto& blockBuf = blockData.block; + uncompressData.blockBuf.MemCopy(blockData.offset, blockData.block.data, blockData.block.curLen); + for (int m = blockData.bamOffset, n = 0; n < blockData.bams.Size(); ++m, ++n) { + auto& bam = blockData.bams.arr[n]; + bam.offset += blockData.offset; + uncompressData.bamArr.arr[m] = bam; + } + } + blockNum += readData.startAddrArr.size(); + } + } +} + +bool doPhase2Uncompress(Phase2PipelineArg& p) { + // spdlog::info("phase2 uncompress order: {}", p.uncompressOrder); + PROF_G_BEG(phase2_uncompress); + bool hasUncompress = false; + int blockNum = 0; + for (int i = 0; i < p.midFiles.size(); ++i) { + auto& f = p.midFiles[i]; + f.needUncompress = false; + if (f.readyReadBufNum > 0 && f.readyUncompressBufNum < Phase2File::UNCOMPRESSS_BUF_NUM) { // 还有空余缓冲区 + blockNum += f.readData[f.uncompressOrder % Phase2File::UNCOMPRESSS_BUF_NUM].startAddrArr.size(); + f.needUncompress = true; + } + } + p.curBlockNum = blockNum; + +#if 1 + kt_for(p.numThread, mtUncompressBlockBatch, &p, p.numThread); + calculateOffset(p); + kt_for(p.numThread, mtCopyToUncompressBuf, &p, p.numThread); +#endif + + for (int i = 0; i < p.midFiles.size(); ++i) { + auto& f = p.midFiles[i]; + if (f.needUncompress) { // 还有空余缓冲区 + f.readyReadBufNum -= 1; + f.readyUncompressBufNum += 1; + f.uncompressOrder += 1; + f.uncompressBuf->startIdx = 0; // 刚读入的，全部bam都还没拷贝到merge + + hasUncompress = true; + } + + // for test + // f.readyUncompressBufNum = 0; + } + PROF_G_END(phase2_uncompress); + return hasUncompress; + + // exit(0); +} + +void* phase2Uncompress(void* data) { + Phase2PipelineArg& p = *(Phase2PipelineArg*)data; + /* 2. do the work */ + while (true) { + // previous dependency + yarn::DEPENDENCY_NOT_TO_BE(p.readSig, 0); + yarn::DEPENDENCY_NOT_TO_BE(p.uncompressSig, Phase2File::UNCOMPRESSS_BUF_NUM); + + if (p.readFinish) { + while (p.uncompressOrder < p.readOrder) { + yarn::DEPENDENCY_NOT_TO_BE(p.uncompressSig, Phase2File::UNCOMPRESSS_BUF_NUM); + doPhase2Uncompress(p); + yarn::UPDATE_SIG_ORDER(p.uncompressSig, p.uncompressOrder); + } + yarn::SIGNAL_FINISH(p.uncompressSig, p.uncompressFinish); + break; + } + doPhase2Uncompress(p); + // if (doPhase2Uncompress(p)) { // 应该不需要判断？ + // update status + yarn::CONSUME_SIGNAL(p.readSig); + yarn::UPDATE_SIG_ORDER(p.uncompressSig, p.uncompressOrder); + // } + } + + spdlog::info("phase2 uncompress end order: {}", p.uncompressOrder); + return nullptr; +} + +//////////////////////////////////////////////// 拷贝到merge的缓冲区 + +int doPhase2CopyToMerge(Phase2PipelineArg& p) { + PROF_G_BEG(phase2_copyToMerge); + int usedUncompress = 0; + p.uncompressReadyNum = 0; +#if 1 + for (int i = 0; i < p.midFiles.size(); ++i) { + int usedUncompressInFile = 0; + auto& f = p.midFiles[i]; + if (f.readyUncompressBufNum > 0 && f.readyMergeBufNum < f.COPY_BUF_NUM) { // 有数据 + // 拷贝到merge数据里的循环数组和循环缓冲区 + auto& uncompressBuf = f.uncompressBuf[f.copyOrder % f.UNCOMPRESSS_BUF_NUM]; + auto& mergeData = f.mergeData; + + if (!mergeData.initialized) { + mergeData.InitSize(uncompressBuf.blockBuf.curLen, uncompressBuf.bamArr.Size()); + } + + int copiedNum = mergeData.CopyBams(uncompressBuf.blockBuf, uncompressBuf.bamArr, uncompressBuf.startIdx); + uncompressBuf.startIdx += copiedNum; + if (uncompressBuf.Size() == 0) { + f.readyUncompressBufNum -= 1; + usedUncompressInFile += 1; + f.copyOrder += 1; + // 检查还能不能继续添加数据 + if (f.readyUncompressBufNum > 0) { + auto& uncompressBuf = f.uncompressBuf[f.copyOrder % f.UNCOMPRESSS_BUF_NUM]; + if (mergeData.hasSpace(uncompressBuf.Front()->wholeBamLen)) { // merge有空间，uncompress有数据，那就继续添加到merge + int copiedNum = mergeData.CopyBams(uncompressBuf.blockBuf, uncompressBuf.bamArr, uncompressBuf.startIdx); + uncompressBuf.startIdx += copiedNum; + if (uncompressBuf.Size() == 0) { + usedUncompressInFile += 1; + f.readyUncompressBufNum -= 1; + f.copyOrder += 1; + } + } + } + } + f.readyMergeBufNum += 1; // 当前文件的merge data数据准备好了 + + //spdlog::info("copy num: {}", copiedNum); + +#if 1 + // for test，消耗mergedata + //f.mergeData.bams.Clear(); + //f.mergeData.data.Clear(); + int clearSize = mergeData.Size() / 3; + //int clearSize = mergeData.Size(); + for (int j = 0; j < clearSize; ++j) { + mergeData.Pop(); + } + f.readyMergeBufNum -= 1; +#endif + } + + usedUncompress = MAX(usedUncompressInFile, usedUncompress); + p.uncompressReadyNum += f.readyUncompressBufNum; + } +#else + p.midFiles[p.copyMergeOrder % p.midFiles.size()].readyUncompressBufNum -= 1; +#endif + spdlog::info("copy to merge order: {}-{}", p.copyMergeOrder, usedUncompress); + + PROF_G_END(phase2_copyToMerge); + + return usedUncompress; +} + +void* phase2CopyToMergeBuf(void* data) { + Phase2PipelineArg& p = *(Phase2PipelineArg*)data; + /* 2. do the work */ + while (true) { + // previous dependency + yarn::DEPENDENCY_NOT_TO_BE(p.uncompressSig, 0); + yarn::DEPENDENCY_NOT_TO_BE(p.copyMerge, Phase2File::COPY_BUF_NUM); + + if (p.uncompressFinish) { + while (p.uncompressReadyNum > 0) { // 一直有没处理完的解压数据 + yarn::DEPENDENCY_NOT_TO_BE(p.copyMerge, Phase2File::COPY_BUF_NUM); + doPhase2CopyToMerge(p); + yarn::UPDATE_SIG_ORDER(p.copyMerge, p.copyMergeOrder); + } + yarn::SIGNAL_FINISH(p.copyMerge, p.copyMergeFinish); + break; + } + int usedUncompress = doPhase2CopyToMerge(p); + + // update status + yarn::CONSUME_SIGNAL_BY(p.uncompressSig, usedUncompress); + //yarn::CONSUME_SIGNAL(p.uncompressSig); + yarn::UPDATE_SIG_ORDER(p.copyMerge, p.copyMergeOrder); + } + + spdlog::info("phase2 copyMerge order: {}", p.copyMergeOrder); + return nullptr; +} \ No newline at end of file diff --git a/src/sort/phase_2_read.h b/src/sort/phase_2_read.h index e69de29..5098835 100644 --- a/src/sort/phase_2_read.h +++ b/src/sort/phase_2_read.h @@ -0,0 +1,16 @@ +/* + Description: 中间文件读入和解析等 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/05 +*/ + +#pragma once + +void* phase2ReadMidFile(void* data); + +void* phase2Uncompress(void* data); + +void* phase2CopyToMergeBuf(void* data); \ No newline at end of file diff --git a/src/sort/phase_2_write.cpp b/src/sort/phase_2_write.cpp index e69de29..714a8d2 100644 --- a/src/sort/phase_2_write.cpp +++ b/src/sort/phase_2_write.cpp @@ -0,0 +1,106 @@ +/* + Description: 写入最终排序后的文件 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/05 +*/ + +#include "phase_2_write.h" + +#include +#include + +#include "phase_2.h" +#include "util/profiling.h" + +//////////////////////////////压缩 +static void mtCompressBlock(void* data, long idx, int tid) { + Phase2PipelineArg& p = *(Phase2PipelineArg*)data; + auto& mergeData = p.mergeData[p.compressOrder % p.COMPRESS_BUF_NUM]; + auto& blockData = mergeData.blockDataArr[idx]; + auto& compressedData = mergeData.compressDataArr[idx]; + + compressedData.ReAllocMem(SINGLE_BLOCK_SIZE); // 压缩后的block数据不会超过单个block的大小 + compressedData.curLen = SINGLE_BLOCK_SIZE; + bgzfCompress(compressedData.data, &compressedData.curLen, blockData.blockBuf.data, blockData.blockBuf.curLen, p.compressLevel); +} + +static void doCompress(Phase2PipelineArg& p) { + PROF_G_BEG(compress); + auto& compressBuf = p.compressBuf[p.compressOrder % p.COMPRESS_BUF_NUM]; + auto& mergeData = p.mergeData[p.compressOrder % p.MERGE_BUF_NUM]; + compressBuf.Clear(); + + kt_for(p.numThread, mtCompressBlock, &p, mergeData.blockDataArr.size()); + + for (int i = 0; i < mergeData.blockDataArr.size(); ++i) { + compressBuf.MemCopy(mergeData.compressDataArr[i].data, mergeData.compressDataArr[i].curLen); + } + mergeData.Clear(); + // spdlog::info("compress bytes: {}", compressBuf.curLen); + PROF_G_END(compress); +} + +/* phase1Compress step- 压缩线程 */ +void* phase2Compress(void* data) { + Phase2PipelineArg& p = *(Phase2PipelineArg*)data; + /* do the work */ + while (true) { + // previous dependency + yarn::DEPENDENCY_NOT_TO_BE(p.mergeSig, 0); + yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, p.COMPRESS_BUF_NUM); + + if (p.mergeFinish) { + while (p.compressOrder < p.mergeOrder) { + yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, p.COMPRESS_BUF_NUM); + doCompress(p); + yarn::UPDATE_SIG_ORDER(p.compressSig, p.compressOrder); + } + yarn::SIGNAL_FINISH(p.compressSig, p.compressFinish); + break; + } + + doCompress(p); + + // update status + yarn::CONSUME_SIGNAL(p.mergeSig); + yarn::UPDATE_SIG_ORDER(p.compressSig, p.compressOrder); + } + + spdlog::info("End phase2 compress order: {}", p.compressOrder); + return nullptr; +} + +/////////////////////////////// 写文件 +static void doPhase2Write(Phase2PipelineArg& p) { + PROF_G_BEG(write_final); + DataBuffer& compressBuf = p.compressBuf[p.writeOrder % p.COMPRESS_BUF_NUM]; + fwrite(compressBuf.data, 1, compressBuf.curLen, p.outFilePtr); + PROF_G_END(write_final); +} + +void* phase2Write(void* data) { + Phase2PipelineArg& p = *(Phase2PipelineArg*)data; + /* do the work */ + while (true) { + // previous dependency + yarn::DEPENDENCY_NOT_TO_BE(p.compressSig, 0); + + if (p.compressFinish) { + while (p.writeOrder < p.compressOrder) { + doPhase2Write(p); + p.writeOrder += 1; + } + break; + } + doPhase2Write(p); + // update status + yarn::CONSUME_SIGNAL(p.compressSig); + p.writeOrder += 1; + } + + spdlog::info("End phase2 write order: {}", p.writeOrder); + return nullptr; +} \ No newline at end of file diff --git a/src/sort/phase_2_write.h b/src/sort/phase_2_write.h index e69de29..b361b02 100644 --- a/src/sort/phase_2_write.h +++ b/src/sort/phase_2_write.h @@ -0,0 +1,14 @@ +/* + Description: 写入最终排序后的文件 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/05 +*/ + +#pragma once + +void* phase2Compress(void* data); + +void* phase2Write(void* data); \ No newline at end of file diff --git a/src/sort/sam_io.cpp b/src/sort/sam_io.cpp index 2efc87d..3e80499 100644 --- a/src/sort/sam_io.cpp +++ b/src/sort/sam_io.cpp @@ -1,12 +1,22 @@ +/* + Description: bam压缩，解压和读入头文件等函数 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/05 +*/ + #include "sam_io.h" #include #include +#include #include #include #include -#include +#include "common_data.h" #include "const_val.h" namespace nsgv { @@ -135,6 +145,11 @@ size_t readUncompressOneBlock(FILE *fpr, uint8_t *fBuf, DataBuffer *uDataPtr) { readState = fread(fBuf, 1, BLOCK_HEADER_LENGTH, fpr); // 先读取一个gz block的头部 blockLen = unpackInt16(&fBuf[16]) + 1; // block的字节数 readState = fread(&fBuf[BLOCK_HEADER_LENGTH], 1, blockLen - BLOCK_HEADER_LENGTH, fpr); + + // 把压缩后的数据记录一下 + nsgv::gInHdr.compressed.ReAllocMem(blockLen); + nsgv::gInHdr.compressed.MemCopy(fBuf, blockLen); + // 解压gz block uint32_t crc = le_to_u32(fBuf + blockLen - 8); size_t newDataSize = uData.maxLen; diff --git a/src/sort/sam_io.h b/src/sort/sam_io.h index 0fe36f0..ab693a9 100644 --- a/src/sort/sam_io.h +++ b/src/sort/sam_io.h @@ -1,3 +1,12 @@ +/* + Description: bam压缩，解压和读入头文件等函数 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/05 +*/ + #pragma once #include #include @@ -23,8 +32,14 @@ struct DataBuffer { data = (uint8_t *)malloc(maxLen); } ~DataBuffer() { - if (data) + if (data) { free(data); + data = nullptr; + } + readPos = 0; + lastPos = 0; + curLen = 0; + maxLen = 0; } void AllocMem(size_t memSize) { curLen = 0; @@ -47,6 +62,15 @@ struct DataBuffer { } } + void ReSize(size_t memSize) { + ReAllocMem(memSize); + curLen = memSize; + } + + void MemCopy(size_t offset, uint8_t* src, size_t len) { + memcpy(&data[offset], src, len); + } + void Clear() { curLen = 0; readPos = 0; @@ -58,6 +82,7 @@ struct HeaderBuf { uint8_t *data = nullptr; // 保留除header之外的解压数据（即最开始的bam记录） int dataLen = 0; // 解压数据的字节数 sam_hdr_t *header; + DataBuffer compressed; }; static int unpackInt16(const uint8_t *buffer) { return buffer[0] | buffer[1] << 8; } diff --git a/src/sort/sort.cpp b/src/sort/sort.cpp index 93e14f9..3e4d7fe 100644 --- a/src/sort/sort.cpp +++ b/src/sort/sort.cpp @@ -1,3 +1,12 @@ +/* + Description: 排序入口 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/05 +*/ + #include "sort.h" #include @@ -42,14 +51,16 @@ static void samSortFirstPipe() { static void bamSortPipeline() { /* set up phase 1 parameters */ Phase1PipelineArg p1; - p1.numThread = nsgv::gSortArg.NUM_THREADS; + p1.SetThreadNum(nsgv::gSortArg.NUM_THREADS); // 平均每线程4M缓冲区，累加起来，用来读入文件（BAM/SAM）（相对解压之后的缓冲区，大小可以忽略） const size_t kReadBufSize = 4L * 1024 * 1024 * p1.numThread; // p1.midFileNamePrefix = getFileNameWithoutExt(nsgv::gSortArg.OUTPUT_FILE) + "_sort_mid."; p1.midFileNamePrefix = getFilePathWithoutExt(nsgv::gSortArg.OUTPUT_FILE) + "_sort_mid."; // 跟输出文件路径相同 spdlog::info("prefix: {}", p1.midFileNamePrefix); p1.maxMemBytes = nsgv::gSortArg.MAX_MEM; - p1.uncompressBufBytes = nsgv::gSortArg.MAX_MEM; // 可以再考虑一下 + + p1.uncompressBufBytes = p1.maxMemBytes; // 可以再考虑一下 + for (int i = 0; i < p1.UNCOMPRESS_BUF_NUM; ++i) { p1.threadUncompressWrap[i].Resize(p1.numThread); // 每个线程的解压block数组初始大小，后续如果不够用会自动扩容 } @@ -69,11 +80,20 @@ static void bamSortPipeline() { ///////////////////////////////////////////////////// // 第二阶段参数初始化 - Phase2PipelineArg p2; + Phase2PipelineArg p2(p1.uncompressData, p1.allBams, p1.midFileNamePrefix, p1.midFileOrder, p1.numThread); + p2.outFilePtr = fopen(nsgv::gSortArg.OUTPUT_FILE.c_str(), "wb"); + if (p2.outFilePtr == nullptr) { + spdlog::error("failed to create output file: {}", nsgv::gSortArg.OUTPUT_FILE.c_str()); + exit(1); + } + // 处理bam的header + fwrite(nsgv::gInHdr.compressed.data, 1, nsgv::gInHdr.compressed.curLen, p2.outFilePtr); // 运行第二阶段 phase2Pipeline(p2); - + // 写结尾 + fwrite("\037\213\010\4\0\0\0\0\0\377\6\0\102\103\2\0\033\0\3\0\0\0\0\0\0\0\0\0", 1, 28, p2.outFilePtr); + fclose(p2.outFilePtr); } // 排序的入口函数，entry function diff --git a/src/sort/sort.h b/src/sort/sort.h index 2a34982..153d0b7 100644 --- a/src/sort/sort.h +++ b/src/sort/sort.h @@ -1,3 +1,12 @@ +/* + Description: 排序相关的数据结构 + + Copyright : All right reserved by ICT + + Author : Zhang Zhonghai + Date : 2026/06/05 +*/ + #pragma once #include @@ -12,8 +21,13 @@ using std::priority_queue; using std::vector; -//#define START_IDX(i, nt, nele) ((i) * (nele) / (nt)) -//#define STOP_IDX(i, nt, nele) (((i)+1) * (nele) / (nt)) +#define START_IDX(i, nt, nele) ((i) * (nele) / (nt)) +#define STOP_IDX(i, nt, nele) (((i) + 1) * (nele) / (nt)) + +#undef MAX +#undef MIN +#define MAX(x, y)((x) > (y) ? (x) : (y)) +#define MIN(x, y)((x) < (y) ? (x) : (y)) /* for step-1 read data from bam file */ struct ReadBuffer { @@ -28,8 +42,16 @@ struct ReadBuffer { blockBuf = (uint8_t *)malloc(SINGLE_BLOCK_SIZE); } ~ReadBuffer() { - if (dataBuf) free(dataBuf); - if (blockBuf) free(blockBuf); + if (dataBuf) { + free(dataBuf); + dataBuf = nullptr; + } + if (blockBuf) { + free(blockBuf); + blockBuf = nullptr; + } + readBufSize = 0; + startAddrArr.clear(); } void Resize(size_t readBufSize_) { if (dataBuf) free(dataBuf); @@ -56,13 +78,27 @@ struct UncompressBlockBuffer { dataBuf = (uint8_t *)malloc(dataBufSize); } ~UncompressBlockBuffer() { - if (dataBuf) free(dataBuf); + if (dataBuf) { + free(dataBuf); + dataBuf = nullptr; + } + dataBufSize = 0; + usedBufSize = 0; + lastEndPos = 0; } void Resize(uint64_t dataBufSize_) { if (dataBuf) free(dataBuf); dataBufSize = dataBufSize_; dataBuf = (uint8_t *)malloc(dataBufSize); } + + void ReAllocate(uint64_t dataBufSize_) { + if (dataBufSize < dataBufSize_) { + dataBufSize = dataBufSize_; + dataBuf = (uint8_t*)realloc(dataBuf, dataBufSize); + } + } + void Clear() { usedBufSize = 0; lastEndPos = 0; @@ -101,6 +137,16 @@ struct FastVector { #endif } } + + T& Get(size_t idx) { return arr[idx]; } + + T* Back() { + if (curIdx == 0) { + return nullptr; + } + return &arr[curIdx - 1]; + } + void Add(const T& item) { T &newItem = Add(); newItem = item; @@ -115,6 +161,20 @@ struct FastVector { arr.resize(num); } } + + void ReSize(size_t num) { + ReAllocate(num); + curIdx = num; + } + + void ClearResize(size_t num) { + if (num > arr.size()) { + arr.clear(); + arr.resize(num); + } + curIdx = num; + } + size_t Size() const { return curIdx; } size_t Capacity() const { return arr.size(); } void Clear() { curIdx = 0; } @@ -127,8 +187,11 @@ struct OneBam { uint32_t wholeBamLen = 0; // 包含4字节bam长度的所有bam内容长度，uint16只适合二代 int32_t tid = 0; // 比对到的染色体 int64_t pos = 0; // mapping 位置 + uint8_t *addr = nullptr; // 首地址 uint64_t offset = 0; // 距离首地址的偏移量 }; typedef FastVector BamArr; -typedef FastVector BamPtrArr; \ No newline at end of file +typedef FastVector BamPtrArr; + +extern void CheckBam(uint8_t* addr, int len); \ No newline at end of file diff --git a/src/sort/sort_args.h b/src/sort/sort_args.h index 669c94f..88c7bfe 100644 --- a/src/sort/sort_args.h +++ b/src/sort/sort_args.h @@ -1,5 +1,5 @@ /* -Description: Markduplicate需要用到的一些参数 +Description: 排序需要用到的一些参数 Copyright : All right reserved by ICT diff --git a/src/util/profiling.cpp b/src/util/profiling.cpp index 43b9f92..d80d0a6 100644 --- a/src/util/profiling.cpp +++ b/src/util/profiling.cpp @@ -63,9 +63,14 @@ int displayProfiling(int nthread) { PRINT_GP(merge); PRINT_GP(compress); PRINT_GP(write_mid); + PRINT_GP(phase2_compress); PRINT_GP(read_mid); PRINT_GP(write_final); PRINT_GP(mid_all); + PRINT_GP(phase2_read); + PRINT_GP(phase2_uncompress); + PRINT_GP(phase2_copyToMerge); + PRINT_GP(phase2); PRINT_TP(sort, nthread); PRINT_TP(mem_copy, nthread); diff --git a/src/util/profiling.h b/src/util/profiling.h index 8cd87bc..197359a 100644 --- a/src/util/profiling.h +++ b/src/util/profiling.h @@ -76,8 +76,14 @@ enum { GP_merge, GP_compress, GP_write_mid, + GP_phase2_read, + GP_phase2_uncompress, + GP_phase2_copyToMerge, + GP_phase2_merge, + GP_phase2_compress, GP_read_mid, - GP_write_final + GP_write_final, + GP_phase2, }; // THREAD enum { TP_0 = 0, TP_1, TP_2, TP_3, TP_4, TP_5, TP_6, TP_7, TP_8, TP_9, TP_10 }; diff --git a/src/util/yarn.h b/src/util/yarn.h index 094925c..3265663 100644 --- a/src/util/yarn.h +++ b/src/util/yarn.h @@ -157,6 +157,10 @@ void free_lock_(lock_t *, char const *, long); order += 1; \ twist_(sig, yarn::BY, 1, __FILE__, __LINE__); +#define UPDATE_SIG(sig) \ + possess_(sig, __FILE__, __LINE__); \ + twist_(sig, yarn::BY, 1, __FILE__, __LINE__); + #define SIGNAL_FINISH(sig, finish) \ possess_(sig, __FILE__, __LINE__); \ finish = 1; \ @@ -166,6 +170,10 @@ void free_lock_(lock_t *, char const *, long); possess_(sig, __FILE__, __LINE__); \ twist_(sig, yarn::BY, -1, __FILE__, __LINE__); +#define CONSUME_SIGNAL_BY(sig, num) \ + possess_(sig, __FILE__, __LINE__); \ + twist_(sig, yarn::BY, num, __FILE__, __LINE__); + #define INIT_SIG(sig) \ possess_(sig, __FILE__, __LINE__); \ twist_(sig, yarn::TO, 0, __FILE__, __LINE__);