fix: some bugs in ivf index

src/core/algorithm/ivf/ivf_params.h

@@ -70,6 +70,8 @@ static const std::string PARAM_IVF_SEARCHER_OPTIMIZER_PARAMS(
     "proxima.ivf.searcher.optimizer_params");
 static const std::string PARAM_IVF_SEARCHER_CONVERTER_REFORMER(
     "proxima.ivf.searcher.converter_reformer");
+static const std::string PARAM_IVF_SEARCHER_NPROBE(
+    "proxima.ivf.searcher.nprobe");
 
 // Constants
 static constexpr char const *kIPMetricName = "InnerProduct";

src/core/algorithm/ivf/ivf_searcher_context.h

@@ -71,9 +71,30 @@ class IVFSearcherContext : public IndexSearcher::Context {
         std::max(static_cast<uint32_t>(
                      std::round(entity_->inverted_list_count() * scan_ratio_)),
                  1u);
+
+    uint32_t nprobe = 0;
+    params.get(PARAM_IVF_SEARCHER_NPROBE, &nprobe);
+    if (nprobe > 0) {
+      nprobe = std::min(nprobe,
+                        static_cast<uint32_t>(entity_->inverted_list_count()));
+      topk_val = nprobe;
+    }
+
     centroid_searcher_ctx_->set_topk(topk_val);
-    max_scan_count_ =
-        static_cast<uint32_t>(std::ceil(entity_->vector_count() * scan_ratio_));
+
+    // When nprobe is explicitly set, scale max_scan_count proportionally
+    // to ensure all probed clusters can be fully scanned.
+    if (nprobe > 0 && entity_->inverted_list_count() > 0) {
+      uint32_t list_count =
+          static_cast<uint32_t>(entity_->inverted_list_count());
+      max_scan_count_ = static_cast<uint32_t>(
+          (static_cast<uint64_t>(entity_->vector_count()) * nprobe +
+           list_count - 1) /
+          list_count);
+    } else {
+      max_scan_count_ = static_cast<uint32_t>(
+          std::ceil(entity_->vector_count() * scan_ratio_));
+    }
     max_scan_count_ = std::max(bruteforce_threshold_, max_scan_count_);
     return 0;
   }

src/core/interface/indexes/ivf_index.cc

@@ -211,9 +211,7 @@ int IVFIndex::_prepare_for_search(
   if (ivf_search_param->nprobe > 0) {
     // TODO: 1. sparse; 2. default ef
     ailego::Params params;
-    // need fix
-    params.set(core::PARAM_IVF_BUILDER_CENTROID_COUNT,
-               ivf_search_param->nprobe);
+    params.set(core::PARAM_IVF_SEARCHER_NPROBE, ivf_search_param->nprobe);
     context->update(params);
   }
   return 0;

tests/ailego/io/file_test.cc

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <cstring>
+#include <vector>
 #include <ailego/utility/memory_helper.h>
 #include <gtest/gtest.h>
 #include <zvec/ailego/io/file.h>
@@ -338,3 +340,65 @@ TEST(File, Seek) {
   EXPECT_EQ((ssize_t)file.size(), file.offset());
   file.close();
 }
+
+TEST(File, MemoryWarmup) {
+  // Test with null address - should not crash
+  File::MemoryWarmup(nullptr, 1024);
+
+  // Test with zero length - should not crash
+  char buf[1];
+  File::MemoryWarmup(buf, 0);
+
+  // Test with a small buffer
+  {
+    const size_t size = 4096;
+    std::vector<uint8_t> buffer(size, 0xAA);
+    File::MemoryWarmup(buffer.data(), buffer.size());
+    // Verify buffer is still intact after warmup
+    EXPECT_EQ(buffer[0], 0xAA);
+    EXPECT_EQ(buffer[size - 1], 0xAA);
+  }
+
+  // Test with a large buffer spanning multiple pages
+  {
+    const size_t size = MemoryHelper::PageSize() * 16 + 100;
+    std::vector<uint8_t> buffer(size, 0x55);
+    File::MemoryWarmup(buffer.data(), buffer.size());
+    EXPECT_EQ(buffer[0], 0x55);
+    EXPECT_EQ(buffer[size - 1], 0x55);
+  }
+
+  // Test with memory-mapped file
+  {
+    const char *file_path = "file_warmup_testing.tmp";
+    size_t file_size = MemoryHelper::PageSize() * 4;
+
+    File file;
+    ASSERT_TRUE(file.create(file_path, file_size));
+
+    void *addr = file.map(0, file_size, File::MMAP_SHARED);
+    ASSERT_TRUE(addr != nullptr);
+
+    // Fill with data
+    memset(addr, 0x77, file_size);
+    EXPECT_TRUE(File::MemoryFlush(addr, file_size));
+    File::MemoryUnmap(addr, file_size);
+    file.close();
+
+    // Open read-only and warmup
+    ASSERT_TRUE(file.open(file_path, true));
+    addr = file.map(0, file_size, 0);
+    ASSERT_TRUE(addr != nullptr);
+
+    File::MemoryWarmup(addr, file_size);
+
+    // Verify data is accessible after warmup
+    uint8_t *p = reinterpret_cast<uint8_t *>(addr);
+    EXPECT_EQ(p[0], 0x77);
+    EXPECT_EQ(p[file_size - 1], 0x77);
+
+    File::MemoryUnmap(addr, file_size);
+    file.close();
+    File::Delete(file_path);
+  }
+}

tests/core/algorithm/ivf/ivf_builder_test.cc

@@ -404,6 +404,70 @@ TEST_F(IVFBuilderTest, TestTrainClusterParams) {
   EXPECT_EQ(0, builder.dump(dumper));
 }
 
+TEST_F(IVFBuilderTest, TestBuildWithConverterClass) {
+  IVFBuilder builder;
+  Params params;
+  params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4");
+  params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster");
+  params.set(PARAM_IVF_BUILDER_CONVERTER_CLASS, "HalfFloatConverter");
+
+  int ret = builder.init(index_meta_, params);
+  EXPECT_EQ(0, ret);
+
+  prepare_index_holder(0, 1000);
+
+  ret = builder.train(threads_, holder_);
+  EXPECT_EQ(0, ret);
+
+  ret = builder.build(threads_, holder_);
+  EXPECT_EQ(0, ret);
+
+  auto centroid_index = builder.centroid_index();
+  EXPECT_GT(centroid_index->centroids_count(), 0u);
+
+  IndexDumper::Pointer dumper = IndexFactory::CreateDumper("MemoryDumper");
+  ret = dumper->create("path");
+  EXPECT_EQ(0, ret);
+
+  ret = builder.dump(dumper);
+  EXPECT_EQ((size_t)1000, builder.stats().built_count());
+  EXPECT_EQ((size_t)1000, builder.stats().dumped_count());
+  EXPECT_EQ((size_t)0, builder.stats().discarded_count());
+}
+
+TEST_F(IVFBuilderTest, TestBuildWithConverterClassMultiLevel) {
+  IVFBuilder builder;
+  Params params;
+  params.set(PARAM_IVF_BUILDER_CENTROID_COUNT, "4*2");
+  params.set(PARAM_IVF_BUILDER_CLUSTER_CLASS, "KmeansCluster*KmeansCluster");
+  params.set(PARAM_IVF_BUILDER_CONVERTER_CLASS, "HalfFloatConverter");
+
+  int ret = builder.init(index_meta_, params);
+  EXPECT_EQ(0, ret);
+
+  prepare_index_holder(0, 1000);
+
+  ret = builder.train(threads_, holder_);
+  EXPECT_EQ(0, ret);
+
+  ret = builder.build(threads_, holder_);
+  EXPECT_EQ(0, ret);
+
+  auto centroid_index = builder.centroid_index();
+  EXPECT_EQ(centroid_index->centroids_count(), 8);
+
+  IndexDumper::Pointer dumper = IndexFactory::CreateDumper("FileDumper");
+  ret = dumper->create("./ivf_converter_test.index");
+  EXPECT_EQ(0, ret);
+
+  ret = builder.dump(dumper);
+  EXPECT_EQ((size_t)1000, builder.stats().built_count());
+  EXPECT_EQ((size_t)1000, builder.stats().dumped_count());
+  EXPECT_EQ((size_t)0, builder.stats().discarded_count());
+  EXPECT_EQ(0, dumper->close());
+  File::RemovePath("./ivf_converter_test.index");
+}
+
 TEST_F(IVFBuilderTest, TestIndexThreads) {
   IndexBuilder::Pointer builder1 = IndexFactory::CreateBuilder("IVFBuilder");
   ASSERT_NE(builder1, nullptr);