K-way merge join

In a previous article, we saw how a branchless algorithm can improve the speed of a two-way merge join. Let’s now explore how we can apply the same principles to a K-way merge join.

A brief recap🔗

Let’s first review what we are trying to achieve here: a join takes two or more datasets and “stitches” them together based on a subset of their columns, called the join key (which can be composite). In this article, we are only interested in equi-joins, where the keys must match exactly. A merge join takes advantage of the datasets being already sorted according to the join key. (Technically, you can merge-join unsorted data by sorting it first, but nobody does that because it is prohibitively expensive: the cost of sorting usually dominates the cost of joining.) Finally, a K-way join merges $K$ lanes of data instead of just two.

Benchmark setup🔗

We will be using our favorite dinosaur for this task, i.e. C++. The code is available online; look at join_bench.cpp in particular.

We try to reproduce a realistic workload: the input is made of $K$ lanes of (columnar) data, and we output the matching indices of each lane (again, in columnar format). This focuses the benchmark on the join operation itself, without having to materialize the joined data. This approach of lazy materialization is more typical of column-oriented than row-oriented computation kernels, although it can be used in both.

Each lane is made of random data from the same domain. We tweak the sparsity so that each input value has an approximate 50% chance of being output. (We do an inner join, so all lanes must match.)

Heap-based version🔗

The canonical implementation of a K-way sorted merge uses a heap to keep track of which lane is currently the minimum. This approach is naturally suited to a plain sorted merge (as used in a merge sort) because it needs to keep duplicates, so that we can just repeatedly pop the smallest lane and keep going. For a merge join, however, things are little trickier, as we need to remove duplicates. That means we need to pop as many values as matching lanes at each step.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

template <typename T>
std::vector<std::vector<size_t>> join_merge_inner_kway_heap(
    const std::span<const std::span<const T>>& inputs)
{
    const size_t k = inputs.size();
    if (k < 2)
        throw std::logic_error("A join requires at least two inputs");

    std::vector<std::vector<size_t>> selected_indices;
    selected_indices.resize(k);

    // Compute number of non-empty lanes
    size_t kk = k;
    for (size_t j = 0; j < k; ++j) {
        if (inputs[j].empty()) {
            kk -= 1;
        }
    }
    if (kk < k)
        return selected_indices;

    // Initialize heap
    std::vector<T> heap;
    heap.reserve(k);
    for (size_t j = 0; j < k; ++j) {
        heap.push_back(inputs[j][0]);
    }
    std::make_heap(heap.begin(), heap.end(), std::greater<>());

    std::vector<size_t> indices;
    indices.resize(k, 0);
    std::vector<bool> lane_matches;
    lane_matches.resize(k);
    while (true) {
        // Peek minimum value
        auto min_value = heap[0];

        // Determine matching lanes
        bool all_matched = true;
        for (size_t j = 0; j < k; ++j) {
            const bool matches = inputs[j][indices[j]] == min_value;
            lane_matches[j] = matches;
            all_matched &= matches;
        }

        if (all_matched) {
            // Output matching indices
            for (size_t j = 0; j < k; ++j) {
                selected_indices[j].push_back(indices[j]);
            }
        }

        // Increment iterators of lanes matching the minimum.
        // At the same time, pop their value from the heap.
        // (It doesn't matter in what order we pop them, as long as we pop the right count.)
        // Also, add the next value if the lane isn't empty.
        bool any_end_reached = false;
        for (size_t j = 0; j < k; ++j) {
            if (lane_matches[j]) {
                indices[j] += 1;
                std::pop_heap(heap.begin(), heap.end(), std::greater<>());
                heap.pop_back();
                if (indices[j] == inputs[j].size()) {
                    any_end_reached = true;
                } else {
                    heap.push_back(inputs[j][indices[j]]);
                    std::push_heap(heap.begin(), heap.end(), std::greater<>());
                }
            }
        }

        if (any_end_reached)
            break;
    }

    return selected_indices;
}

With this version, we achieve 101M items/s with 4 lanes, down to 23.8M items/s with 1024 lanes.

Branchless version🔗

Can we do better? To answer this question, let’s try to identify our limiting factor. Clearly, data locality is good, as we read each lane sequentially, and the heap has a bounded size. A potential problem, however, is branch mispredictions: we register 153k of them for 4 lanes. The cause is easy to spot: a heap is an inherently “branchy” data structure, as popping requires reorganizing it, which causes unpredictable branches and can involve an arbitrary number of layers.

That’s why, instead of using a heap to maintain the current minimum, we can recompute the minimum at each step. This approach is somewhat counter-intuitive, as it should use more instructions. Also, computing a minimum is prone to data dependencies, although we saw a while ago that compilers know how to speed up this kind of operation through auto-vectorization. So, maybe it’s not as expensive as it looks?

In addition, we apply the same principles as we did for the two-way merge join: instead of conditionally outputting matching indices, we always output the current indices, but only advance the output pointer when the join keys matched. That requires reserving extra space for the output indices, and trimming the result at the end.

The result is a little convoluted, so bear with me: 🐻

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

template <typename T>
std::vector<std::vector<size_t>> join_merge_inner_kway_branchless(
    const std::span<const std::span<const T>>& inputs)
{
    const size_t k = inputs.size();
    if (k < 2)
        throw std::logic_error("A join requires at least two inputs");

    std::vector<std::vector<size_t>> selected_indices;
    selected_indices.resize(k);

    std::vector<size_t> indices;
    indices.resize(k, 0);
    size_t oi = 0;
    const size_t buffer_size = 128;

    // Compute number of non-empty lanes
    size_t kk = k;
    for (size_t j = 0; j < k; ++j) {
        if (inputs[j].empty()) {
            kk -= 1;
        }
    }
    if (kk < k)
        return selected_indices;

    std::vector<bool> lane_matches;
    lane_matches.resize(k);
    while (true) {
        // Resize output vectors so that we can assign values instead of pushing them
        if (oi >= selected_indices[0].size()) {
            for (size_t j = 0; j < k; ++j)
                selected_indices[j].resize(oi + buffer_size);
        }

        // Compute minimum value across all lanes
        auto min_value = inputs[0][indices[0]];
        for (size_t j = 1; j < k; ++j) {
            min_value = std::min(min_value, inputs[j][indices[j]]);
        }

        // Unconditionnally store indices in the output vectors
        for (size_t j = 0; j < k; ++j)
            selected_indices[j][oi] = indices[j];

        // Determine if all lanes matched
        bool all_matched = true;
        for (size_t j = 0; j < k; ++j) {
            const bool matches = inputs[j][indices[j]] == min_value;
            lane_matches[j] = matches;
            all_matched &= matches;
        }

        // Only increment the output index if LHS and RHS matched.
        // That way, if they don't, the stored values will be discarded
        oi += all_matched;

        // Increment iterators of lanes matching the minimum
        for (size_t j = 0; j < k; ++j) {
            indices[j] += lane_matches[j];
        }

        // Determine if any end reached
        bool any_end_reached = false;
        for (size_t j = 0; j < k; ++j) {
            any_end_reached |= indices[j] == inputs[j].size();
        }
        if (any_end_reached)
            break;
    }

    // Resize selected indices to drop the extraneous buffer
    for (size_t j = 0; j < k; ++j)
        selected_indices[j].resize(oi);

    return selected_indices;
}

This new version achieves 151M items/s for 4 lanes, down to 81.4M/s for 1024 lanes. This is consistently better than the heap-based version, and if we except the case with four lanes, more than twice as fast!

Can the speedup be attributed to branch mispredictions? Not entirely. The branchless code does exhibit fewer branch misses… but only when many lanes are involved.

However, it consistently uses fewer instructions. This is counter-intuitive, as we would expect to do more work by always recomputing the minimum from scratch, instead of relying on the heap to work its magic.

Maybe that’s due to auto-vectorization? The branchless version is made of multiple inner loops, in most of which iterations are independent—the exceptions being computing the minimum, and determining if all lanes matched. And even those can be partially unrolled.

Alas, again, our intuition is baffled: when forbidding automating unrolling and vectorization, we do get marginally higher instruction counts, but still consistenly fewer than the heap-based version. Also, the throughput remains unchanged. (Sorry, no hard numbers for this; you will just need to take my word for it…)

What’s happening, then? The crux of the matter is that we are doing a join, which implies that we expect at least some of the rows to match. When they do, the heap-based version needs to pop as many times as there are matching values, whereas the branchless version just increments matching pointers. Therefore, although a heap is the most efficient way of maintaining a minimum across iterations, the branchless version “cheats” by advancing as many lanes as possible.

Conclusion🔗

The K-way merge join behaves just like its two-way variant: branchless code can surpass branchy code, not only because it causes fewer branch mispredictions, but also because its simplicity allows it to use fewer instructions. There is a kind of “compounding effect” at work here.

Our little benchmark also shows that heaps, despite being very smart data structures with optimal theoretical complexity, are heavily branchy, full of data dependencies, and therefore ill-suited to modern CPUs. (Which doesn’t mean, of course, that they can’t be useful in targeted situations, like Heapselect.)