SegmentBuilder test

CMakeLists.txt

@@ -95,8 +95,8 @@ set(fpindexlib_SOURCES
 	src/fpindex/segment_file_format.h
 	src/fpindex/segment.h
 	src/fpindex/segment.cpp
-	src/fpindex/in_memory_segment.h
-	src/fpindex/in_memory_segment.cpp
+	src/fpindex/segment_builder.h
+	src/fpindex/segment_builder.cpp
 	src/fpindex/io/memory_file.h
 	src/fpindex/io/memory_file.cpp
 	src/fpindex/proto/internal.pb.h
@@ -236,6 +236,7 @@ set(tests_SOURCES
 	src/index/segment_merge_policy_test.cpp
 	src/index/op_test.cpp
 	src/fpindex/segment_test.cpp
+	src/fpindex/segment_builder_test.cpp
 	src/store/buffered_input_stream_test.cpp
 	src/store/input_stream_test.cpp
 	src/store/output_stream_test.cpp

src/fpindex/segment.cpp

@@ -137,4 +137,10 @@ bool Segment::Load(const std::shared_ptr<io::File>& file) {
     return true;
 }
 
+Segment::Segment(uint32_t id, std::shared_ptr<io::File> file, SegmentHeader header, size_t first_block_offset,
+                 std::vector<std::pair<uint32_t, uint32_t>>&& block_index)
+    : BlockBasedSegment(id), header_(header), file_(file), first_block_offset_(first_block_offset), block_index_(std::move(block_index)) {
+        ready_ = true;
+}
+
 }  // namespace fpindex

src/fpindex/segment.h

@@ -40,6 +40,9 @@ class Segment : public BlockBasedSegment {
     bool GetBlock(size_t block_no, std::vector<std::pair<uint32_t, uint32_t>> *items) override;
 
  private:
+    Segment(uint32_t id, std::shared_ptr<io::File> file, SegmentHeader header, size_t first_block_offset,
+            std::vector<std::pair<uint32_t, uint32_t>> &&block_index);
+
     std::mutex mutex_;
     std::atomic<bool> ready_{false};
     SegmentHeader header_;

src/fpindex/in_memory_segment.cpp → src/fpindex/segment_builder.cpp

@@ -1,4 +1,4 @@
-#include "fpindex/in_memory_segment.h"
+#include "fpindex/segment_builder.h"
 
 #include <google/protobuf/io/coded_stream.h>
 
@@ -9,7 +9,7 @@
 
 namespace fpindex {
 
-bool InMemorySegment::Add(uint32_t id, const std::vector<uint32_t>& hashes) {
+bool SegmentBuilder::Add(uint32_t id, const std::vector<uint32_t>& hashes) {
     std::unique_lock<std::shared_mutex> lock(mutex_);
     if (frozen_) {
         return false;
@@ -20,14 +20,19 @@ bool InMemorySegment::Add(uint32_t id, const std::vector<uint32_t>& hashes) {
     return true;
 }
 
-void InMemorySegment::Freeze() {
+bool SegmentBuilder::IsFrozen() {
+    std::unique_lock<std::shared_mutex> lock(mutex_);
+    return frozen_;
+}
+
+void SegmentBuilder::Freeze() {
     if (!frozen_) {
         std::unique_lock<std::shared_mutex> lock(mutex_);
         frozen_ = true;
     }
 }
 
-bool InMemorySegment::Search(const std::vector<uint32_t>& hashes, std::vector<SearchResult>* results) {
+bool SegmentBuilder::Search(const std::vector<uint32_t>& hashes, std::vector<SearchResult>* results) {
     std::shared_lock<std::shared_mutex> lock;
     if (!frozen_) {
         lock = std::shared_lock<std::shared_mutex>(mutex_);
@@ -47,7 +52,7 @@ bool InMemorySegment::Search(const std::vector<uint32_t>& hashes, std::vector<Se
     return true;
 }
 
-bool InMemorySegment::Serialize(io::File* file) {
+bool SegmentBuilder::Serialize(io::File* file) {
     std::shared_lock<std::shared_mutex> lock;
     if (!frozen_) {
         lock = std::shared_lock<std::shared_mutex>(mutex_);

src/fpindex/in_memory_segment.h → src/fpindex/segment_builder.h

@@ -10,18 +10,19 @@
 
 namespace fpindex {
 
-class InMemorySegment : public BaseSegment {
+class SegmentBuilder : public BaseSegment {
  public:
-    InMemorySegment(uint32_t id) : BaseSegment(id) {}
-    InMemorySegment(const InMemorySegment&) = delete;
-    InMemorySegment& operator=(const InMemorySegment&) = delete;
+    SegmentBuilder(uint32_t id) : BaseSegment(id) {}
+    SegmentBuilder(const SegmentBuilder&) = delete;
+    SegmentBuilder& operator=(const SegmentBuilder&) = delete;
 
     bool Add(uint32_t id, const std::vector<uint32_t>& values);
 
     bool Search(const std::vector<uint32_t>& hashes, std::vector<SearchResult>* results) override;
 
     // Freeze the segment so that no more data can be added to it.
     void Freeze();
+    bool IsFrozen();
 
     // Serialize the segment data to the output stream.
     bool Serialize(io::File* file);

src/fpindex/segment_builder_test.cpp

@@ -0,0 +1,27 @@
+#include "fpindex/segment_builder.h"
+
+#include <gtest/gtest.h>
+
+using namespace fpindex;
+using namespace testing;
+
+TEST(SegmentBuilderTest, SearchEmpty) {
+    SegmentBuilder segment(0);
+
+    std::vector<uint32_t> query{1, 2, 3};
+    std::vector<SearchResult> results;
+    ASSERT_TRUE(segment.Search(query, &results));
+    ASSERT_TRUE(results.empty());
+}
+
+TEST(SegmentBuilderTest, SearchExactMatch) {
+    SegmentBuilder segment(0);
+    segment.Add(1, {1, 2, 3});
+
+    std::vector<uint32_t> query{1, 2, 3};
+    std::vector<SearchResult> results;
+    ASSERT_TRUE(segment.Search(query, &results));
+    ASSERT_EQ(1, results.size());
+    ASSERT_EQ(1, results[0].id());
+    ASSERT_EQ(3, results[0].score());
+}

src/fpindex/segment_file_format.h

@@ -88,7 +88,7 @@ inline Iter SerializeSegmentBlock(io::CodedOutputStream *output, const SegmentHe
     return end;
 }
 
-bool ParseSegmentBlockV1(io::CodedInputStream *input, const SegmentHeader &header,
+inline bool ParseSegmentBlockV1(io::CodedInputStream *input, const SegmentHeader &header,
                          std::vector<std::pair<uint32_t, uint32_t>> *entries) {
     const int block_size = header.block_size();
     auto start = input->CurrentPosition();
@@ -128,7 +128,7 @@ bool ParseSegmentBlockV1(io::CodedInputStream *input, const SegmentHeader &heade
     return true;
 }
 
-bool ParseSegmentBlock(io::CodedInputStream *input, const SegmentHeader &header,
+inline bool ParseSegmentBlock(io::CodedInputStream *input, const SegmentHeader &header,
                        std::vector<std::pair<uint32_t, uint32_t>> *entries) {
     if (header.block_format() == SegmentBlockFormat::BLOCK_FORMAT_V1) {
         return ParseSegmentBlockV1(input, header, entries);