Faster in-place union of a few containers into many

src/roaring.c

@@ -1064,6 +1064,76 @@ static void roaring_inplace_merge_bulk(roaring_bitmap_t *x1,
     *ra1 = merged;
 }
 
+// Union of x2 into x1 when x2 has few containers compared to x1. We find each
+// key of x2 by galloping, and we insert the new containers by moving the
+// containers of x1 once, block by block, from the back.
+static void roaring_inplace_union_few(roaring_bitmap_t *x1,
+                                      const roaring_bitmap_t *x2) {
+    roaring_array_t *ra1 = &x1->high_low_container;
+    const roaring_array_t *ra2 = &x2->high_low_container;
+    const bool cow2 = is_cow(x2);
+    const int32_t length1 = ra1->size;
+    const int32_t length2 = ra2->size;
+    uint8_t result_type = 0;
+
+    // Union of the containers with a key in both, counting the other ones.
+    int32_t added = 0;
+    int32_t pos1 = -1;
+    for (int32_t pos2 = 0; pos2 < length2; pos2++) {
+        const uint16_t key = ra2->keys[pos2];
+        pos1 = ra_advance_until(ra1, key, pos1);
+        if (pos1 == length1 || ra1->keys[pos1] != key) {
+            added++;
+            pos1--;  // The next key may be at the same position.
+            continue;
+        }
+        uint8_t type1 = ra1->typecodes[pos1];
+        container_t *c1 = ra1->containers[pos1];
+        if (container_is_full(c1, type1)) continue;
+        container_t *c2 = ra2->containers[pos2];
+        uint8_t type2 = ra2->typecodes[pos2];
+        container_t *c =
+            (type1 == SHARED_CONTAINER_TYPE)
+                ? container_or(c1, type1, c2, type2, &result_type)
+                : container_ior(c1, type1, c2, type2, &result_type);
+        if (c != c1) {
+            container_free(c1, type1);
+        }
+        ra_set_container_at_index(ra1, pos1, c, result_type);
+    }
+    if (added == 0) return;
+
+    // Insertion of the other containers, from the back: the containers of x1
+    // from index 'end' on are already in place.
+    extend_array(ra1, added);
+    int32_t end = length1;
+    int32_t dst = length1 + added;
+    for (int32_t pos2 = length2 - 1; pos2 >= 0 && dst > end; pos2--) {
+        const uint16_t key = ra2->keys[pos2];
+        int32_t i = binarySearch(ra1->keys, end, key);
+        if (i >= 0) continue;
+        i = -i - 1;
+        const int32_t n = end - i;
+        dst -= n;
+        memmove(ra1->keys + dst, ra1->keys + i, n * sizeof(uint16_t));
+        memmove(ra1->containers + dst, ra1->containers + i,
+                n * sizeof(container_t *));
+        memmove(ra1->typecodes + dst, ra1->typecodes + i, n * sizeof(uint8_t));
+        dst--;
+        uint8_t type2 = ra2->typecodes[pos2];
+        container_t *c2 =
+            get_copy_of_container(ra2->containers[pos2], &type2, cow2);
+        if (cow2) {
+            ra_set_container_at_index(ra2, pos2, c2, type2);
+        }
+        ra1->keys[dst] = key;
+        ra1->containers[dst] = c2;
+        ra1->typecodes[dst] = type2;
+        end = i;
+    }
+    ra1->size = length1 + added;
+}
+
 // inplace or (modifies its first argument).
 void roaring_bitmap_or_inplace(roaring_bitmap_t *x1,
                                const roaring_bitmap_t *x2) {
@@ -1077,6 +1147,11 @@ void roaring_bitmap_or_inplace(roaring_bitmap_t *x1,
         roaring_bitmap_overwrite(x1, x2);
         return;
     }
+    // Merging would visit all the containers of x1.
+    if ((int64_t)length2 * 32 < length1) {
+        roaring_inplace_union_few(x1, x2);
+        return;
+    }
     int pos1 = 0, pos2 = 0;
     uint8_t type1, type2;
     uint16_t s1 = ra_get_key_at_index(&x1->high_low_container, (uint16_t)pos1);

tests/toplevel_unit.c

@@ -3876,6 +3876,74 @@ DEFINE_TEST(test_subset) {
     roaring_bitmap_free(rb2);
 }
 
+// In-place union of a bitmap with few containers into one with many.
+DEFINE_TEST(test_or_inplace_few_containers) {
+    uint32_t seed = 1234;
+    for (int trial = 0; trial < 200; trial++) {
+        // Mixing bitmaps with and without copy-on-write is unsupported.
+        bool cow1 = trial & 1, cow2 = cow1;
+        roaring_bitmap_t *big = roaring_bitmap_create();
+        roaring_bitmap_set_copy_on_write(big, cow1);
+        // Keys 2k, with array, bitset, run and full containers.
+        for (uint32_t k = 0; k < 1000; k++) {
+            uint32_t base = (2 * k) << 16;
+            switch (k % 4) {
+                case 0:
+                    roaring_bitmap_add(big, base + k);
+                    break;
+                case 1:
+                    for (uint32_t v = 0; v < 65536; v += 3)
+                        roaring_bitmap_add(big, base + v);
+                    break;
+                case 2:
+                    roaring_bitmap_add_range(big, base + 100, base + 5000);
+                    break;
+                default:
+                    roaring_bitmap_add_range(big, base, base + 65536);
+            }
+        }
+        roaring_bitmap_run_optimize(big);
+        roaring_bitmap_t *small = roaring_bitmap_create();
+        roaring_bitmap_set_copy_on_write(small, cow2);
+        int nkeys = 1 + trial % 20;
+        for (int j = 0; j < nkeys; j++) {
+            seed = seed * 1103515245 + 12345;
+            // Existing keys, keys between them, before them and after them.
+            uint32_t key = (seed >> 8) % 2004;
+            if (j == 0 && trial % 3 == 0) key = 2003;
+            uint32_t base = key << 16;
+            if (seed & 1) {
+                roaring_bitmap_add_range(small, base + (seed >> 20),
+                                         base + 65536);
+            } else {
+                roaring_bitmap_add(small, base + (seed >> 16));
+            }
+        }
+        roaring_bitmap_t *shared = NULL;
+        if (cow1) shared = roaring_bitmap_copy(big);  // Shares containers.
+        roaring_bitmap_t *expected = roaring_bitmap_or(big, small);
+        roaring_bitmap_t *small_copy = roaring_bitmap_copy(small);
+        roaring_bitmap_or_inplace(big, small);
+        const char *reason = NULL;
+        assert_true(roaring_bitmap_internal_validate(big, &reason));
+        assert_true(roaring_bitmap_equals(big, expected));
+        assert_true(roaring_bitmap_equals(small, small_copy));
+        if (shared) {
+            roaring_bitmap_t *expected2 = roaring_bitmap_or(shared, small);
+            roaring_bitmap_free(shared);
+            assert_true(roaring_bitmap_equals(big, expected2));
+            roaring_bitmap_free(expected2);
+        }
+        roaring_bitmap_free(small);
+        // The containers copied from 'small' must outlive it.
+        assert_true(roaring_bitmap_internal_validate(big, &reason));
+        assert_true(roaring_bitmap_equals(big, expected));
+        roaring_bitmap_free(small_copy);
+        roaring_bitmap_free(expected);
+        roaring_bitmap_free(big);
+    }
+}
+
 DEFINE_TEST(test_or_many_memory_leak) {
     for (int i = 0; i < 10; i++) {
         roaring_bitmap_t *bm1 = roaring_bitmap_create();
@@ -5739,6 +5807,7 @@ int main() {
     tellmeall();
 
     const struct CMUnitTest tests[] = {
+        cmocka_unit_test(test_or_inplace_few_containers),
         cmocka_unit_test(issue743),
         cmocka_unit_test(fuzz_deserializer),
         cmocka_unit_test(issue660),