Make external sort based offline index building work (infiniflow#812)

### What problem does this PR solve?

Previously, offline index building works the same as the near real time
index building which is either time inefficient and could take much
memory.

Issue link:infiniflow#633

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Test cases

Author: yingfeng

src/storage/invertedindex/column_inverter.cpp

@@ -19,7 +19,6 @@ module;
 #include <cstdio>
 #include <cstring>
 #include <vector>
-
 module column_inverter;
 import stl;
 import analyzer;
@@ -200,32 +199,37 @@ void ColumnInverter::GeneratePosting() {
     }
 }
 
+void ColumnInverter::SortForOfflineDump() {
+    MergePrepare();
+    Sort();
+}
+
 /// Layout of the input of external sort file
-//    +-----------++----------------++--------------------++--------------------------++-------------------------------------------------------+
-//    |           ||                ||                    ||                          ||                                                       |
-//    |   Count   ||  Size of A Run ||   Num of records   ||   Position of Next Run   ||             Data of a Run                             |
-//    |           ||                ||   within a Run     ||                          ||                                                       |
-//    +-----------++----------------++--------------------++--------------------------++-------------------------------------------------------+
-//                 ----------------------------------------------------------------------------------------------------------------------------+
+//    +-----------+  +----------------++--------------------++--------------------------++-------------------------------------------------------+
+//    |           |  |                ||                    ||                          ||                                                       |
+//    |   Count   |  |  Size of A Run ||   Num of records   ||   Position of Next Run   ||             Data of a Run                             |
+//    |           |  |                ||   within a Run     ||                          ||                                                       |
+//    +-----------+  +----------------++--------------------++--------------------------++-------------------------------------------------------+
+//                   ----------------------------------------------------------------------------------------------------------------------------+
 //                                                            Data within each group
 void ColumnInverter::SpillSortResults(FILE *spill_file, u64 &tuple_count) {
     // spill sort results for external merge sort
+
+    // size of this Run in bytes
     u32 data_size = 0;
     u32 data_size_pos = ftell(spill_file);
     fwrite(&data_size, sizeof(u32), 1, spill_file);
     // number of tuples
     u32 num_of_tuples = positions_.size();
     tuple_count += num_of_tuples;
     fwrite(&num_of_tuples, sizeof(u32), 1, spill_file);
-
     // start offset for next spill
     u64 next_start_offset = 0;
     u64 next_start_offset_pos = ftell(spill_file);
     fwrite(&next_start_offset, sizeof(u64), 1, spill_file);
-
     u32 data_start_offset = ftell(spill_file);
     // sorted data
-    u32 last_term_num = 0;
+    u32 last_term_num = std::numeric_limits<u32>::max();
     StringRef term;
     u16 record_length = 0;
     char str_null = '\0';
@@ -246,9 +250,7 @@ void ColumnInverter::SpillSortResults(FILE *spill_file, u64 &tuple_count) {
     next_start_offset = ftell(spill_file);
     data_size = next_start_offset - data_start_offset;
     fseek(spill_file, data_size_pos, SEEK_SET);
-    fwrite(&data_size, sizeof(u32), 1, spill_file);
-
-    // update offset for next spill
+    fwrite(&data_size, sizeof(u32), 1, spill_file); // update offset for next spill
     fseek(spill_file, next_start_offset_pos, SEEK_SET);
     fwrite(&next_start_offset, sizeof(u64), 1, spill_file);
     fseek(spill_file, next_start_offset, SEEK_SET);

src/storage/invertedindex/column_inverter.cppm

@@ -43,6 +43,8 @@ public:
 
     void InvertColumn(u32 doc_id, const String &val);
 
+    void SortForOfflineDump();
+
     void Merge(ColumnInverter &rhs);
 
     static void Merge(Vector<SharedPtr<ColumnInverter>> &inverters);

src/storage/invertedindex/common/external_sort_merger.cppm

@@ -53,6 +53,7 @@ export struct TermTuple {
     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_)
@@ -74,9 +75,9 @@ export struct TermTuple {
 
     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; }
 
-    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));

src/storage/invertedindex/index_defines.cppm

@@ -49,8 +49,8 @@ export {
     constexpr u32 MAX_UNCOMPRESSED_SKIP_LIST_SIZE = 10;
     constexpr u8 SKIP_LIST_BUFFER_SIZE = 32;
 
-    constexpr const char *DICT_SUFFIX = "dic";
-    constexpr const char *POSTING_SUFFIX = "pos";
+    constexpr const char *DICT_SUFFIX = ".dic";
+    constexpr const char *POSTING_SUFFIX = ".pos";
 
     using ScoredId = Pair<float, u32>;
     using ScoredIds = Vector<ScoredId>;

src/storage/invertedindex/memory_indexer.cpp

@@ -78,7 +78,7 @@ MemoryIndexer::~MemoryIndexer() {
     Reset();
 }
 
-void MemoryIndexer::Insert(SharedPtr<ColumnVector> column_vector, u32 row_offset, u32 row_count) {
+void MemoryIndexer::Insert(SharedPtr<ColumnVector> column_vector, u32 row_offset, u32 row_count, bool offline) {
     u64 seq_inserted(0);
     u32 doc_count(0);
     {
@@ -90,16 +90,38 @@ void MemoryIndexer::Insert(SharedPtr<ColumnVector> column_vector, u32 row_offset
     }
     auto task = MakeShared<BatchInvertTask>(seq_inserted, column_vector, row_offset, row_count, doc_count);
     PostingWriterProvider provider = [this](const String &term) -> SharedPtr<PostingWriter> { return GetOrAddPosting(term); };
-    auto func = [this, task, provider](int id) {
-        auto inverter = MakeShared<ColumnInverter>(this->analyzer_, &this->byte_slice_pool_, provider);
-        inverter->InvertColumn(task->column_vector_, task->row_offset_, task->row_count_, task->start_doc_id_);
-        this->ring_inverted_.Put(task->task_seq_, inverter);
-    };
-    thread_pool_.push(func);
+    if (offline) {
+        auto func = [this, task, provider](int id) {
+            auto inverter = MakeShared<ColumnInverter>(this->analyzer_, &this->byte_slice_pool_, provider);
+            inverter->InvertColumn(task->column_vector_, task->row_offset_, task->row_count_, task->start_doc_id_);
+            inverter->SortForOfflineDump();
+            this->ring_sorted_.Put(task->task_seq_, inverter);
+        };
+        thread_pool_.push(func);
+    } else {
+        auto func = [this, task, provider](int id) {
+            auto inverter = MakeShared<ColumnInverter>(this->analyzer_, &this->byte_slice_pool_, provider);
+            inverter->InvertColumn(task->column_vector_, task->row_offset_, task->row_count_, task->start_doc_id_);
+            this->ring_inverted_.Put(task->task_seq_, inverter);
+        };
+        thread_pool_.push(func);
+    }
 }
 
-void MemoryIndexer::Commit() {
-    thread_pool_.push([this](int id) { this->CommitSync(); });
+void MemoryIndexer::Commit(bool offline) {
+    if (offline) {
+        if (nullptr == spill_file_handle_) {
+            PrepareSpillFile();
+        }
+        while (inflight_tasks_ != 0) {
+            this->ring_sorted_.Iterate([this](SharedPtr<ColumnInverter> &inverter) {
+                inverter->SpillSortResults(this->spill_file_handle_, this->tuple_count_);
+                inflight_tasks_--;
+                num_runs_++;
+            });
+        }
+    } else
+        thread_pool_.push([this](int id) { this->CommitSync(); });
 }
 
 SizeT MemoryIndexer::CommitSync() {
@@ -122,12 +144,17 @@ SizeT MemoryIndexer::CommitSync() {
     return num;
 }
 
-void MemoryIndexer::Dump() {
+void MemoryIndexer::Dump(bool offline) {
+    if (offline) {
+        Commit(true);
+        OfflineDump();
+        return;
+    }
+
     while (GetInflightTasks() > 0) {
         sleep(1);
         CommitSync();
     }
-
     Path path = Path(index_dir_) / base_name_;
     String index_prefix = path.string();
     LocalFileSystem fs;
@@ -191,52 +218,53 @@ void MemoryIndexer::OfflineDump() {
     // 2. Generate posting
     // 3. Dump disk segment data
     FinalSpillFile();
+    std::cout << "num_runs_ " << num_runs_ << std::endl;
     SortMerger<TermTuple, u16> *merger = new SortMerger<TermTuple, u16>(spill_full_path_.c_str(), num_runs_, 100000000, 2);
     merger->Run();
     delete merger;
     FILE *f = fopen(spill_full_path_.c_str(), "r");
     u64 count;
     fread((char *)&count, sizeof(u64), 1, f);
-
-    String index_prefix; /// TODO, to be integrated
+    Path path = Path(index_dir_) / base_name_;
+    String index_prefix = path.string();
     LocalFileSystem fs;
     String posting_file = index_prefix + POSTING_SUFFIX;
     SharedPtr<FileWriter> posting_file_writer = MakeShared<FileWriter>(fs, posting_file, 128000);
     String dict_file = index_prefix + DICT_SUFFIX;
     SharedPtr<FileWriter> dict_file_writer = MakeShared<FileWriter>(fs, dict_file, 128000);
     TermMetaDumper term_meta_dumpler((PostingFormatOption(flag_)));
-
     String fst_file = index_prefix + DICT_SUFFIX + ".fst";
     std::ofstream ofs(fst_file.c_str(), std::ios::binary | std::ios::trunc);
     OstreamWriter wtr(ofs);
     FstBuilder builder(wtr);
 
+    u16 record_length;
+    char buf[MAX_TUPLE_LENGTH];
     std::string_view last_term;
     u32 last_term_pos = 0;
     u32 last_doc_id = INVALID_DOCID;
-    u16 record_length;
     Deque<String> term_buffer;
-    char buf[MAX_TUPLE_LENGTH];
     UniquePtr<PostingWriter> posting;
     SizeT term_meta_offset = 0;
+
     for (u64 i = 0; i < count; ++i) {
         fread(&record_length, sizeof(u16), 1, f);
         fread(buf, record_length, 1, f);
         TermTuple tuple(buf, record_length);
         if (tuple.term_ != last_term) {
             if (last_doc_id != INVALID_DOCID) {
                 posting->EndDocument(last_doc_id, 0);
+                if (posting.get()) {
+                    TermMeta term_meta(posting->GetDF(), posting->GetTotalTF());
+                    posting->Dump(posting_file_writer, term_meta);
+                    term_meta_dumpler.Dump(dict_file_writer, term_meta);
+                    builder.Insert((u8 *)last_term.data(), last_term.length(), term_meta_offset);
+                    term_meta_offset = dict_file_writer->TotalWrittenBytes();
+                }
             }
             term_buffer.push_back(String(tuple.term_));
             std::string_view view(term_buffer.back());
             last_term.swap(view);
-            if (posting.get()) {
-                TermMeta term_meta(posting->GetDF(), posting->GetTotalTF());
-                posting->Dump(posting_file_writer, term_meta);
-                term_meta_dumpler.Dump(dict_file_writer, term_meta);
-                builder.Insert((u8 *)last_term.data(), last_term.length(), term_meta_offset);
-                term_meta_offset = dict_file_writer->TotalWrittenBytes();
-            }
             posting = MakeUnique<PostingWriter>(&byte_slice_pool_, &buffer_pool_, PostingFormatOption(flag_));
         } else if (last_doc_id != tuple.doc_id_) {
             posting->EndDocument(last_doc_id, 0);
@@ -247,15 +275,11 @@ void MemoryIndexer::OfflineDump() {
             posting->AddPosition(last_term_pos);
         }
     }
+
     if (last_doc_id != INVALID_DOCID) {
         posting->EndDocument(last_doc_id, 0);
     }
-    if (posting->GetDF() > 0) {
-        TermMeta term_meta(posting->GetDF(), posting->GetTotalTF());
-        posting->Dump(posting_file_writer, term_meta);
-        term_meta_dumpler.Dump(dict_file_writer, term_meta);
-        builder.Insert((u8 *)last_term.data(), last_term.length(), term_meta_offset);
-    }
+    posting_file_writer->Sync();
     dict_file_writer->Sync();
     builder.Finish();
     fs.AppendFile(dict_file, fst_file);

src/storage/invertedindex/memory_indexer.cppm

@@ -51,18 +51,18 @@ public:
     ~MemoryIndexer();
 
     // Insert is non-blocking. Caller must ensure there's no RowID gap between each call.
-    void Insert(SharedPtr<ColumnVector> column_vector, u32 row_offset, u32 row_count);
+    void Insert(SharedPtr<ColumnVector> column_vector, u32 row_offset, u32 row_count, bool offline = false);
 
     // Commit is non-blocking. There shall be a background thread which call this method regularly (for example, every 2 seconds).
     // Other threads can also call this method.
-    void Commit();
+    void Commit(bool offline = false);
 
     // CommitSync wait at max 100ms to get a batch of insertions and commit them. Returens the size of the batch.
     SizeT CommitSync();
 
     // Dump is blocking and shall be called only once after inserting all documents.
     // WARN: Don't reuse MemoryIndexer after calling Dump!
-    void Dump();
+    void Dump(bool offline = false);
 
     SizeT GetInflightTasks() {
         std::unique_lock<std::mutex> lock(mutex_);

src/unit_test/storage/invertedindex/common/external_sort.cpp

@@ -124,13 +124,15 @@ class ExternalSortTest : public BaseTest {
         while (run_num < 100 || SIZE % run_num != 0)
             run_num = rand() % 300;
 
+        run_num = 10;
         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();
+                // str = RandStr();
+                String str = std::to_string(i * SIZE / run_num + j);
                 u32 doc_id = 34567; // i * SIZE / run_num + j;
                 u32 term_pos = i;
                 memcpy(buffer, str.data(), str.size());
@@ -178,4 +180,4 @@ TEST_F(ExternalSortTest, test1) {
     CheckMerger<u32, u8>(10000, 1000000);
 }
 
-TEST_F(ExternalSortTest, test2) { CheckTermTuple(10000, 1000000); }
+TEST_F(ExternalSortTest, test2) { CheckTermTuple(100, 1000000); }

src/unit_test/storage/invertedindex/memory_indexer.cpp

@@ -110,3 +110,51 @@ TEST_F(MemoryIndexerTest, Insert) {
         }
     }
 }
+
+TEST_F(MemoryIndexerTest, test2) {
+    // https://en.wikipedia.org/wiki/Finite-state_transducer
+    const char *paragraphs[] = {
+        R"#(A finite-state transducer (FST) is a finite-state machine with two memory tapes, following the terminology for Turing machines: an input tape and an output tape. This contrasts with an ordinary finite-state automaton, which has a single tape. An FST is a type of finite-state automaton (FSA) that maps between two sets of symbols.[1] An FST is more general than an FSA. An FSA defines a formal language by defining a set of accepted strings, while an FST defines a relation between sets of strings.)#",
+        R"#(An FST will read a set of strings on the input tape and generates a set of relations on the output tape. An FST can be thought of as a translator or relater between strings in a set.)#",
+        R"#(In morphological parsing, an example would be inputting a string of letters into the FST, the FST would then output a string of morphemes.)#",
+        R"#(An automaton can be said to recognize a string if we view the content of its tape as input. In other words, the automaton computes a function that maps strings into the set {0,1}. Alternatively, we can say that an automaton generates strings, which means viewing its tape as an output tape. On this view, the automaton generates a formal language, which is a set of strings. The two views of automata are equivalent: the function that the automaton computes is precisely the indicator function of the set of strings it generates. The class of languages generated by finite automata is known as the class of regular languages.)#",
+        R"#(The two tapes of a transducer are typically viewed as an input tape and an output tape. On this view, a transducer is said to transduce (i.e., translate) the contents of its input tape to its output tape, by accepting a string on its input tape and generating another string on its output tape. It may do so nondeterministically and it may produce more than one output for each input string. A transducer may also produce no output for a given input string, in which case it is said to reject the input. In general, a transducer computes a relation between two formal languages.)#",
+    };
+    const SizeT num_paragraph = sizeof(paragraphs) / sizeof(char *);
+    SharedPtr<ColumnVector> column = ColumnVector::Make(MakeShared<DataType>(LogicalType::kVarchar));
+    column->Initialize();
+    for (SizeT i = 0; i < num_paragraph; ++i) {
+        Value v = Value::MakeVarchar(String(paragraphs[i]));
+        column->AppendValue(v);
+    }
+    Vector<ExpectedPosting> expected_postings = {{"fst", {0, 1, 2}, {4, 2, 2}}, {"automaton", {0, 3}, {2, 5}}, {"transducer", {0, 4}, {1, 4}}};
+
+    MemoryIndexer
+        indexer1("/tmp/infinity/fulltext_tbl1_col1", "chunk1", RowID(0U, 0U), flag_, "standard", byte_slice_pool_, buffer_pool_, thread_pool_);
+    indexer1.Insert(column, 0, 2, true);
+    indexer1.Insert(column, 2, 2, true);
+    indexer1.Insert(column, 4, 1, true);
+    indexer1.Dump(true);
+
+    ColumnIndexReader reader;
+    reader.Open("/tmp/infinity/fulltext_tbl1_col1", Vector<String>{"chunk1"}, Vector<RowID>{RowID(0U, 0U)}, flag_, nullptr);
+    for (SizeT i = 0; i < expected_postings.size(); ++i) {
+        const ExpectedPosting &expected = expected_postings[i];
+        const String &term = expected.term;
+
+        UniquePtr<PostingIterator> post_iter(reader.Lookup(term, &byte_slice_pool_));
+        ASSERT_TRUE(post_iter != nullptr);
+
+        RowID doc_id = INVALID_ROWID;
+        for (SizeT j = 0; j < expected.doc_ids.size(); ++j) {
+            doc_id = post_iter->SeekDoc(expected.doc_ids[j]);
+            ASSERT_EQ(doc_id, expected.doc_ids[j]);
+            u32 tf = post_iter->GetCurrentTF();
+            ASSERT_EQ(tf, expected.tfs[j]);
+        }
+        if (doc_id != INVALID_ROWID) {
+            doc_id = post_iter->SeekDoc(doc_id + 1);
+            ASSERT_EQ(doc_id, INVALID_ROWID);
+        }
+    }
+}