Support string sort in external merge sort (#680)

yingfeng · web-flow · commit 6f8b3cb27b74 · 2024-02-29T16:30:13.000+08:00
### What problem does this PR solve? Improve external merge sort to support tuple[string, uint32, uint32] sorting Issue link:#633 ### Type of change - [x] New Feature (non-breaking change which adds functionality)
diff --git a/src/storage/invertedindex/common/external_sort_merger.cpp b/src/storage/invertedindex/common/external_sort_merger.cpp
@@ -172,10 +172,8 @@ void SortMerger<KeyType, LenType>::Init(DirectIO &io_stream) {
         micro_run_pos_[i] = KeyAddr(micro_buf_[i], -1, i).LEN() + sizeof(LenType);
         num_micro_run_[i] = 0;
 
-        // if size_run_[i]>PRE_BUF_SIZE_, it needs to the end of a run and turn to the next run
         io_stream.Seek(next_run_pos);
     }
-    // LOG_INFO << "Run number: " << group_size_;
 
     // initialize predict heap and records number of every microrun
     for (u32 i = 0; i < group_size_; ++i) {
@@ -409,5 +407,5 @@ void SortMerger<KeyType, LenType>::Run() {
 }
 
 template class SortMerger<u32, u8>;
-
+template class SortMerger<TermTuple, u16>;
 } // namespace infinity
diff --git a/src/storage/invertedindex/common/external_sort_merger.cppm b/src/storage/invertedindex/common/external_sort_merger.cppm
@@ -47,11 +47,149 @@ export struct DirectIO {
     SizeT length_;
 };
 
+export struct TermTuple {
+    std::string_view term_;
+    u32 doc_id_;
+    u32 term_pos_;
+    int Compare(const TermTuple &rhs) const {
+        int ret = term_.compare(rhs.term_);
+        if (ret == 0) {
+            if (doc_id_ != rhs.doc_id_) {
+                if (doc_id_ < rhs.doc_id_)
+                    return -1;
+                else
+                    return 1;
+            } else {
+                if (term_pos_ != rhs.term_pos_) {
+                    if (term_pos_ < rhs.term_pos_)
+                        return -1;
+                    else
+                        return 1;
+                }
+                return 0;
+            }
+        } else
+            return ret < 0 ? -1 : 1;
+    }
+
+    bool operator==(const TermTuple &other) const { return Compare(other) == 0; }
+
+    bool operator>(const TermTuple &other) const { return Compare(other) < 0; }
+
+    bool operator<(const TermTuple &other) const { return Compare(other) > 0; }
+
+    TermTuple(char *p, u16 len) : term_(p, len - sizeof(doc_id_) - sizeof(term_pos_) - 1) {
+        doc_id_ = *((u32 *)(p + term_.size() + 1));
+        term_pos_ = *((u32 *)(p + term_.size() + 1 + sizeof(doc_id_)));
+    }
+};
+
+template <typename KeyType, typename LenType, typename = void>
+struct KeyAddress {
+    char *data{nullptr};
+    u64 addr;
+    u32 idx;
+
+    KeyAddress(char *p, u64 ad, u32 i) {
+        data = p;
+        addr = ad;
+        idx = i;
+    }
+
+    KeyAddress() {
+        data = nullptr;
+        addr = -1;
+        idx = -1;
+    }
+
+    KeyType KEY() { return KeyType(data + sizeof(LenType), LEN()); }
+    KeyType KEY() const { return KeyType(data + sizeof(LenType), LEN()); }
+    LenType LEN() const { return *(LenType *)data; }
+    u64 &ADDR() { return addr; }
+    u64 ADDR() const { return addr; }
+    u32 IDX() const { return idx; }
+    u32 &IDX() { return idx; }
+
+    int Compare(const KeyAddress &p) const {
+        if (KEY() == p.KEY())
+            return 0;
+        else if (KEY() > p.KEY())
+            return 1;
+        else
+            return -1;
+    }
+
+    bool operator==(const KeyAddress &other) const { return Compare(other) == 0; }
+
+    bool operator>(const KeyAddress &other) const { return Compare(other) < 0; }
+
+    bool operator<(const KeyAddress &other) const { return Compare(other) > 0; }
+};
+
+template <typename KeyType, typename LenType>
+struct KeyAddress<KeyType, LenType, typename std::enable_if<std::is_scalar<KeyType>::value>::type> {
+    char *data{nullptr};
+    u64 addr;
+    u32 idx;
+
+    KeyAddress(char *p, u64 ad, u32 i) {
+        data = p;
+        addr = ad;
+        idx = i;
+    }
+
+    KeyAddress() {
+        data = nullptr;
+        addr = -1;
+        idx = -1;
+    }
+
+    KeyType &KEY() { return *(KeyType *)(data + sizeof(LenType)); }
+    KeyType KEY() const { return *(KeyType *)(data + sizeof(LenType)); }
+
+    LenType LEN() const { return *(LenType *)data; }
+    u64 &ADDR() { return addr; }
+    u64 ADDR() const { return addr; }
+    u32 IDX() const { return idx; }
+    u32 &IDX() { return idx; }
+
+    int Compare(const KeyAddress &p) const {
+        if (KEY() == p.KEY())
+            return 0;
+        else if (KEY() > p.KEY())
+            return 1;
+        else
+            return -1;
+
+        LenType len1 = LEN() / sizeof(KeyType);
+        LenType len2 = p.LEN() / sizeof(KeyType);
+
+        for (LenType i = 1; i < len1 && i < len2; ++i) {
+            if (((KeyType *)(data + sizeof(LenType)))[i] > ((KeyType *)(p.data + sizeof(LenType)))[i])
+                return 1;
+            if (((KeyType *)(data + sizeof(LenType)))[i] < ((KeyType *)(p.data + sizeof(LenType)))[i])
+                return -1;
+        }
+
+        if (len1 == len2)
+            return 0;
+
+        if (len1 > len2)
+            return 1;
+        return -1;
+    }
+
+    bool operator==(const KeyAddress &other) const { return Compare(other) == 0; }
+
+    bool operator>(const KeyAddress &other) const { return Compare(other) < 0; }
+
+    bool operator<(const KeyAddress &other) const { return Compare(other) > 0; }
+};
+
 export template <typename KeyType, typename LenType>
 class SortMerger {
     typedef SortMerger<KeyType, LenType> self_t;
-
-    struct KeyAddr;
+    typedef KeyAddress<KeyType, LenType> KeyAddr;
     String filenm_;
     const u32 MAX_GROUP_SIZE_; //!< max group size
     const u32 BS_SIZE_;        //!< in fact it equals to memory size
@@ -60,8 +198,8 @@ class SortMerger {
     u32 OUT_BUF_SIZE_;         //!< max size of output buffer
     const u32 OUT_BUF_NUM_;    //!< output threads number
 
-    std::priority_queue<struct KeyAddr> pre_heap_;   //!< predict heap
-    std::priority_queue<struct KeyAddr> merge_heap_; //!< merge heap
+    std::priority_queue<KeyAddr> pre_heap_;   //!< predict heap
+    std::priority_queue<KeyAddr> merge_heap_; //!< merge heap
 
     u32 *micro_run_idx_{nullptr};   //!< the access index of each microruns
     u32 *micro_run_pos_{nullptr};   //!< the access position within each microruns
@@ -102,65 +240,6 @@ class SortMerger {
 
     u64 FILE_LEN_;
 
-    struct KeyAddr {
-        char *data{nullptr};
-        u64 addr;
-        u32 idx;
-
-        KeyAddr(char *p, u64 ad, u32 i) {
-            data = p;
-            addr = ad;
-            idx = i;
-        }
-
-        KeyAddr() {
-            data = nullptr;
-            addr = -1;
-            idx = -1;
-        }
-
-        KeyType &KEY() { return *(KeyType *)(data + sizeof(LenType)); }
-        KeyType KEY() const { return *(KeyType *)(data + sizeof(LenType)); }
-        LenType LEN() const { return *(LenType *)data; }
-        u64 &ADDR() { return addr; }
-        u64 ADDR() const { return addr; }
-        u32 IDX() const { return idx; }
-        u32 &IDX() { return idx; }
-
-        int Compare(const KeyAddr &p) const {
-            if (KEY() == p.KEY())
-                return 0;
-
-            if (KEY() > p.KEY())
-                return 1;
-            if (KEY() < p.KEY())
-                return -1;
-
-            LenType len1 = LEN() / sizeof(KeyType);
-            LenType len2 = p.LEN() / sizeof(KeyType);
-
-            for (LenType i = 1; i < len1 && i < len2; ++i) {
-                if (((KeyType *)(data + sizeof(LenType)))[i] > ((KeyType *)(p.data + sizeof(LenType)))[i])
-                    return 1;
-                if (((KeyType *)(data + sizeof(LenType)))[i] < ((KeyType *)(p.data + sizeof(LenType)))[i])
-                    return -1;
-            }
-
-            if (len1 == len2)
-                return 0;
-
-            if (len1 > len2)
-                return 1;
-            return -1;
-        }
-
-        bool operator==(const KeyAddr &other) const { return Compare(other) == 0; }
-
-        bool operator>(const KeyAddr &other) const { return Compare(other) < 0; }
-
-        bool operator<(const KeyAddr &other) const { return Compare(other) > 0; }
-    };
-
     void NewBuffer();
 
     void Init(DirectIO &io_stream);
diff --git a/src/unit_test/storage/invertedindex/common/external_sort.cpp b/src/unit_test/storage/invertedindex/common/external_sort.cpp
@@ -16,6 +16,7 @@
 #include <cassert>
 #include <filesystem>
 #include <iostream>
+#include <string.h>
 
 import stl;
 import third_party;
@@ -38,6 +39,17 @@ Vector<char> RandStr(T key) {
     return str;
 }
 
+Vector<char> RandStr() {
+    u32 len = rand() % 100;
+
+    Vector<char> str;
+    str.reserve(len);
+    for (u32 i = 0; i < len; ++i)
+        str.push_back('a' + rand() % 26);
+
+    return str;
+}
+
 class ExternalSortTest : public BaseTest {
 public:
     ExternalSortTest() {}
@@ -58,7 +70,6 @@ class ExternalSortTest : public BaseTest {
         while (run_num < 100 || SIZE % run_num != 0)
             run_num = rand() % 300;
 
-        // run_num = 800;
         for (u32 i = 0; i < run_num; ++i) {
             u64 pos = ftell(f);
             fseek(f, 2 * sizeof(u32) + sizeof(u64), SEEK_CUR);
@@ -69,8 +80,6 @@ class ExternalSortTest : public BaseTest {
                 fwrite(&len, sizeof(LenType), 1, f);
                 fwrite(str.data(), len, 1, f);
                 s += len + sizeof(LenType);
-
-                // cout<<"\rAdd data: "<<(double)(i*SIZE/run_num+j)/SIZE*100.<<"%"<<std::Flush;
             }
             u64 next_run_pos = ftell(f);
             fseek(f, pos, SEEK_SET);
@@ -103,9 +112,70 @@ class ExternalSortTest : public BaseTest {
         }
         std::filesystem::remove("./tt");
     }
+
+    void CheckTermTuple(const u64 SIZE, u32 bs = 100000000) {
+        std::filesystem::remove("./tt");
+
+        Vector<char> str;
+        FILE *f = fopen("./tt", "w+");
+        fwrite(&SIZE, sizeof(u64), 1, f);
+
+        u32 run_num = rand() % 300;
+        while (run_num < 100 || SIZE % run_num != 0)
+            run_num = rand() % 300;
+
+        char buffer[200];
+        for (u32 i = 0; i < run_num; ++i) {
+            u64 pos = ftell(f);
+            fseek(f, 2 * sizeof(u32) + sizeof(u64), SEEK_CUR);
+            u32 s = 0;
+            for (u32 j = 0; j < SIZE / run_num; ++j) {
+                str = RandStr();
+                u32 doc_id = 34567; // i * SIZE / run_num + j;
+                u32 term_pos = i;
+                memcpy(buffer, str.data(), str.size());
+                buffer[str.size()] = '\0';
+                memcpy(buffer + str.size() + 1, &doc_id, sizeof(u32));
+                memcpy(buffer + str.size() + 1 + sizeof(u32), &term_pos, sizeof(u32));
+                u16 len = str.size() + 1 + sizeof(u32) + sizeof(u32);
+                fwrite(&len, sizeof(u16), 1, f);
+                fwrite(buffer, len, 1, f);
+                s += len + sizeof(u16);
+            }
+            u64 next_run_pos = ftell(f);
+            fseek(f, pos, SEEK_SET);
+            fwrite(&s, sizeof(u32), 1, f);
+            s = SIZE / run_num;
+            fwrite(&s, sizeof(u32), 1, f);
+            fwrite(&next_run_pos, sizeof(u64), 1, f);
+            fseek(f, 0, SEEK_END);
+        }
+        fclose(f);
+
+        SortMerger<TermTuple, u16> merger("./tt", run_num, bs, 2);
+        merger.Run();
+
+        f = fopen("./tt", "r");
+        u64 count = 0;
+        u32 doc_id = 34567;
+        fread(&count, sizeof(u64), 1, f);
+        EXPECT_EQ(count, SIZE);
+        for (u32 i = 0; i < count; ++i) {
+            u16 len = 0;
+            fread(&len, sizeof(u16), 1, f);
+            char *buf = new char[len];
+            fread(buf, len, 1, f);
+            TermTuple tuple(buf, len);
+            EXPECT_EQ(tuple.doc_id_, doc_id);
+            delete[] buf;
+        }
+        std::filesystem::remove("./tt");
+    }
 };
 
 TEST_F(ExternalSortTest, test1) {
     CheckMerger<u32, u8>(1000, 100000);
     CheckMerger<u32, u8>(10000, 1000000);
 }
+
+TEST_F(ExternalSortTest, test2) { CheckTermTuple(10000, 1000000); }

Original file line number	Diff line number	Diff line change
`@@ -172,10 +172,8 @@ void SortMerger<KeyType, LenType>::Init(DirectIO &io_stream) {`
`172`	`172`	`micro_run_pos_[i] = KeyAddr(micro_buf_[i], -1, i).LEN() + sizeof(LenType);`
`173`	`173`	`num_micro_run_[i] = 0;`
`174`	`174`
`175`		`- // if size_run_[i]>PRE_BUF_SIZE_, it needs to the end of a run and turn to the next run`
`176`	`175`	`io_stream.Seek(next_run_pos);`
`177`	`176`	`}`
`178`		`- // LOG_INFO << "Run number: " << group_size_;`
`179`	`177`
`180`	`178`	`// initialize predict heap and records number of every microrun`
`181`	`179`	`for (u32 i = 0; i < group_size_; ++i) {`
`@@ -409,5 +407,5 @@ void SortMerger<KeyType, LenType>::Run() {`
`409`	`407`	`}`
`410`	`408`
`411`	`409`	`template class SortMerger<u32, u8>;`
`412`		`-`
	`410`	`+template class SortMerger<TermTuple, u16>;`
`413`	`411`	`} // namespace infinity`