多线程外排序解决大数据排序问题2(最小堆并行k路归并)

转自：AIfred

事实证明外排序的效率主要依赖于磁盘，归并阶段采用K路归并可以显著减少IO量，最小堆并行k路归并，效率倍增。

二路归并的思路会导致非常多冗余的磁盘访问，两组两组合并确定的是当前的相对位置并不能一次确定最终的位置。

K路归并，每一轮归并直接确定的是最终的位置，不用重复访问，减少IO。该排序算法需要对每个整数做2次磁盘读和2次磁盘写。

摘自维基百科：

外排序的一个例子是外归并排序（External merge sort），它读入一些能放在内存内的数据量，在内存中排序后输出为一个顺串（即是内部数据有序的临时文件），处理完所有的数据后再进行归并。比如，要对900MB 的数据进行排序，但机器上只有100 MB的可用内存时，外归并排序按如下方法操作：

1. 读入100 MB的数据至内存中，用某种常规方式（如快速排序、堆排序、归并排序等方法）在内存中完成排序。
2. 将排序完成的数据写入磁盘。
3. 重复步骤1和2直到所有的数据都存入了不同的100 MB的块（临时文件）中。在这个例子中，有900 MB数据，单个临时文件大小为100 MB，所以会产生9个临时文件。
4. 读入每个临时文件（顺串）的前10 MB（ = 100 MB / (9块 + 1)）的数据放入内存中的输入缓冲区，最后的10 MB作为输出缓冲区。（实践中，将输入缓冲适当调小，而适当增大输出缓冲区能获得更好的效果。）
5. 执行九路归并算法，将结果输出到输出缓冲区。一旦输出缓冲区满，将缓冲区中的数据写出至目标文件，清空缓冲区。一旦9个输入缓冲区中的一个变空，就从这个缓冲区关联的文件，读入下一个10M数据，除非这个文件已读完。这是“外归并排序”能在主存外完成排序的关键步骤 -- 因为“归并算法”(merge algorithm)对每一个大块只是顺序地做一轮访问(进行归并)，每个大块不用完全载入主存。

算法思路：

1. 二分文件位置，选取每一个文件的枢轴，将每一个文件划分为thread个片段，使得每一个thread处理所有文件片段和相对均衡。

2. 然后用每一个线程各自处理属于他们的K个文件片段，规模为K的最小堆维护K路归并，构造一个大小为k的堆，先将k个节点的头元素插入到堆中，然后每次取出头结点，取出来的元素属于哪个子数组，再添加这个子数组的下一个元素进入堆中，来维护这个堆。这里的排序结果也是最后的排序结果，直接输出到文件。多线程并行处理。

用到知识点：

1. 数据并行拆分: （partition_and_sort）

• 切分的大小符合内存大小限制。
• 禁止拆分数据线程间的依赖。
• 汇总时处理并发冲突，原子操作。

2. k路归并堆排序：（heapsort）

构造一个大小为k的最小二叉堆，先将k个节点的头元素插入到堆中，然后每次取出头结点，取出来的元素属于哪个子数组，再添加这个子数组的下一个元素进入堆中，来维护这个堆。这里的排序结果也是最后的排序结果，减少IO操作

3. 缓冲区（buffer）

如果将每个节点的最小值放入内存，例如2，1，3，4放入内存，但是把最小值1拿掉之后需要补充一个元素，将外部内存的2拿到内存里来，可是外部内存可能在硬盘或网络，此过程相比内存操作会慢很多，不断读取外部内存效率很低，所以采用缓冲区，每次读取k个节点前部分数据到内存缓冲区（几k或几M）。

#include <stdio.h>
#include <cstring>
#include <string>
#include <atomic>
#include <queue>
#include <vector>
#include <Windows.h>
#include <ppl.h>
#include <functional>
#include <io.h>
#include <time.h>
#define MAX_THREADS 4
#define MAX_K 100
using namespace std;
using namespace concurrency;
const int dx = 20;//并行快速排序的dx优化
const long long PARTITION_SIZE = 100000000;
const long long BUFFER_SIZE = 200000;
const long long EACH_NUM = (PARTITION_SIZE / MAX_THREADS);
int parts, heapsize[MAX_THREADS];
long long data_size;
mutex m;
typedef pair<int, int> node; // (int，文件id)
// 每一个线程维护的最小堆，堆大小是文件数，K路归并
node heap[MAX_THREADS][MAX_K + 10]; 

void parallel_qsort(int *begin, int *end) {//并行快速排序
    if (begin >= end - 1) return;
    int *key = rand() % (end - begin) + begin;
    swap(*key, *begin);
    int *i = begin, *j = begin;
    for (key = begin; j < end; j++) {
        if (*j < *key) {
            i++;
            swap(*i, *j);
        }
    }
    swap(*begin, *i);
    if (i - begin > dx && end - i > dx) {//dx优化
        parallel_for(0, 2, [&](int x) {
            if (x) parallel_qsort(begin, i);
            else parallel_qsort(i + 1, end);
        });
    } else {
        parallel_qsort(begin, i);
        parallel_qsort(i + 1, end);
    }
}

// 添加新元素，向上找到合适的插入位置
inline void up(int idx, int x) {
    int fa = x >> 1; node tmp = heap[idx][x];
    while (fa) {
        if (tmp < heap[idx][fa])//cmp
            heap[idx][x] = heap[idx][fa];
        else break;
        x = fa; fa = x >> 1;
    }
    heap[idx][x] = tmp;
}

// 向下找到合适的插入位置
inline void down(int idx, int x) {
    int ch = x << 1; node tmp = heap[idx][x];
    while (ch <= heapsize[idx]) {
        if (ch < heapsize[idx] && heap[idx][ch + 1] < heap[idx][ch]) ch++;//cmp
        if (heap[idx][ch] < tmp)//cmp
            heap[idx][x] = heap[idx][ch];
        else break;
        x = ch; ch = x << 1;
    }
    heap[idx][x] = tmp;
}

inline void push(int idx, node val) { // 向最小堆插入元素
    heap[idx][++heapsize[idx]] = val;
    up(idx, heapsize[idx]);
}
inline node top(int idx) { return heap[idx][1]; }
inline void pop(int idx) { // pop堆顶最小元素
    heap[idx][1] = heap[idx][heapsize[idx]--];
    down(idx, 1);
}

inline void ch_size(string file_name, fpos_t size) {
    FILE *fout = fopen(file_name.c_str(), "wb");
    _chsize_s(fileno(fout), size * sizeof(int));
    fclose(fout);
}
inline int seek_dat(FILE* &f, fpos_t pos) {
    int *get = new int;
    pos *= sizeof(int);
    fsetpos(f, &pos);
    fread(get, sizeof(int), 1, f);
    int tmp = *get; delete get;
    return tmp;
}

void partition_and_sort(string in_file) {
    int *arr = new int[PARTITION_SIZE];
    for (long long i = 0; i < (data_size - 1) / PARTITION_SIZE + 1; i++) {
        atomic_int each_get[MAX_THREADS + 1] = {};
        string tmp_file = "temp\part" + to_string(i) + ".dat";
        clock_t start = clock();
        cout << "Reading part " << i << "...";
        parallel_for(0, MAX_THREADS, [&](long long x) {
            FILE* fin = fopen(in_file.c_str(), "rb");
            fpos_t pos = (PARTITION_SIZE * i + EACH_NUM * x) * sizeof(int);
            if (fsetpos(fin, &pos) == 0)
                each_get[x] = fread(arr + EACH_NUM * x, sizeof(int), EACH_NUM, fin);
            each_get[MAX_THREADS] += each_get[x];
            fclose(fin);
        });
        cout << "
Sorting part " << i << "...";
        parallel_qsort(arr, arr + each_get[MAX_THREADS]); // 并行快速排序
        cout << "
Writing part " << i << "...";
        ch_size(tmp_file, each_get[MAX_THREADS]);
        parallel_for(0, MAX_THREADS, [&](long long x) {
            FILE* fout = fopen(tmp_file.c_str(), "rb+");
            fpos_t pos = EACH_NUM * x * sizeof(int);
            if (fsetpos(fout, &pos) == 0)
                fwrite(arr + EACH_NUM * x, sizeof(int), each_get[x], fout);
            fclose(fout);
        });
        clock_t end = clock();
        cout << "
Part " << i << " established. Time usage = " << end - start << "ms.
";
    }
    delete[] arr;
}

void merge_file() {
    FILE* fin[MAX_K] = {};
    fpos_t size[MAX_K] = {}, seek_pos[MAX_THREADS + 1][MAX_K + 1] = {};
    for (int i = 0; i < parts; i++) {
        fin[i] = fopen(("temp\part" + to_string(i) + ".dat").c_str(), "rb");
        fseek(fin[i], 0, SEEK_END);
        fgetpos(fin[i], &size[i]);
        size[i] /= sizeof(int); // 有多少数
        seek_pos[MAX_THREADS][parts] += (seek_pos[MAX_THREADS][i] = size[i]); // seek_pos[线程id][文件id] = 文件位置
    }
    cout << "
Initializing merging operation...
";
    for (long long i = 1; i < MAX_THREADS; i++) {
        fpos_t l0 = 0, r0 = size[0] - 1;
        while (r0 - l0 > 1) {  // 二分文件0的位置
            seek_pos[i][parts] = seek_pos[i][0] = (l0 + r0) / 2;
            int get0 = seek_dat(fin[0], seek_pos[i][0]);
            for (int idx = 1; idx < parts; idx++) {
                fpos_t l = 0, r = size[idx];
                while (r - l > 0) { // 二分其他文件的位置，找到get0
                    seek_pos[i][idx] = (l + r) / 2;
                    int get = seek_dat(fin[idx], seek_pos[i][idx]);
                    if (get0 <= get) r = seek_pos[i][idx];
                    else l = seek_pos[i][idx] + 1;
                }
                seek_pos[i][parts] += (seek_pos[i][idx] = r);
            }
            // 二分文件0位置的目的是使得分治的较为均衡，所有文件相对片段长度之和接近于 data_size / MAX_THREADS
            if (seek_pos[i][parts] * MAX_THREADS < data_size * i) l0 = seek_pos[i][0] + 1;
            else r0 = seek_pos[i][0] - 1;
        }
    }
    for (int i = 0; i < parts; i++) fclose(fin[i]);
    clock_t start = clock(); atomic_llong all_write = 0;
    parallel_for(0, MAX_THREADS, [&](int x) { // 线程处理外循环
        FILE *fin[MAX_K] = {}, *fout = fopen("ans.dat", "rb+");
        fpos_t fpos = seek_pos[x][parts] * sizeof(int);
        fsetpos(fout, &fpos);
        int **buf = new int*[MAX_K + 1];// 开文件数个buffer，K路归并
        for (int i = 0; i <= MAX_K; i++) buf[i] = new int[BUFFER_SIZE];
        int pos[MAX_K + 1] = {};//buffer pos
        fpos_t all[MAX_K] = {};
        for (int i = 0; i < parts; i++) {
            fin[i] = fopen(("temp\part" + to_string(i) + ".dat").c_str(), "rb");
            fpos = seek_pos[x][i] * sizeof(int);
            fsetpos(fin[i], &fpos);
            all[i] = seek_pos[x + 1][i] - seek_pos[x][i]; // 记录每一个文件一个线程处理的长度
            fread(buf[i], sizeof(int), BUFFER_SIZE, fin[i]);
        }
        for (int i = 0; i < parts; i++) {
            // 向最小堆中读入所有文件属于该线程处理的部分的第一个元素
            push(x, node(buf[i][0], i)); //(线程id，pair(buffer，文件id))
            pos[i] = 1; all[i]--;
        }
        while (heapsize[x]) {
            // buf[parts]: k路归并排好序的缓冲区
            if (pos[parts] == BUFFER_SIZE) {
                fwrite(buf[parts], sizeof(int), BUFFER_SIZE, fout);
                all_write += BUFFER_SIZE;
                if (all_write % 1000000 == 0) {
                    m.lock();
                    cout << "
Start merging... " << (all_write * 100) / data_size
                         << "% completed.";
                    m.unlock();
                }
                pos[parts] = 0;
            }
            int bel = top(x).second;
            buf[parts][pos[parts]++] = top(x).first;
            if (all[bel]) {
                heap[x][1] = node(buf[bel][pos[bel]], bel); down(x, 1);// 该buffer的新元素替换heap的顶部最小元素
                if ((++pos[bel]) == BUFFER_SIZE) {
                    fread(buf[bel], sizeof(int), BUFFER_SIZE, fin[bel]);
                    pos[bel] = 0;
                }
                all[bel]--;
            } else pop(x); // 该文件属于线程x的部分全部处理完了，就直接pop
        }
        fwrite(buf[parts], sizeof(int), pos[parts], fout); // 把余下排好序的buffer写入文件
        cout << "
Start merging... 100% completed.";
        for (int i = 0; i < parts; i++) fclose(fin[i]); fclose(fout);
        for (int i = 0; i < MAX_K; i++) delete[] buf[i]; delete[] buf;
    });
    clock_t end = clock();
    cout << "
Merging finished. Time usage = " << end - start << "ms.
";
}


int main() {
    string in_file;
    cout << "Enter data file name: ";
    cin >> in_file;
    FILE* fin = fopen(in_file.c_str(), "rb");
    if (fin == NULL) {
        cout << "Could not open that file.
";
        main();
    }
    clock_t start_time = clock();
    fseek(fin, 0, SEEK_END);
    fgetpos(fin, &data_size);
    data_size /= sizeof(int);
    parts = (data_size - 1) / PARTITION_SIZE + 1;
    fclose(fin);
    cout << "
Partitioning " << data_size << " elements(int)...
";
    system("mkdir temp");
    parallel_for(0, 2, [&](int x) {
        if (x) partition_and_sort(in_file);
        else ch_size("ans.dat", data_size);
    });
    merge_file();
    clock_t end_time = clock();
    system("rd /s/q temp");
    cout << "
External sorting complete, result saved to "ans.dat".
"
         << "Time usage = " << end_time - start_time << "ms.
";
    system("pause");
    return 0;
}

参考：

简单无堆无缓冲区单线程K路归并版本: https://www.cnblogs.com/this-543273659/archive/2011/07/30/2122083.html

wiki:外排序