refactor(quantization, python_bindings): Reconstruct Thread-Local Storage and Python Binding Interfaces

1. Thread-Local Storage Optimization
   - Replaced the thread_local static member `query_` in QuantizerBase with a local static variable inside a function
   -彻底 resolved symbol duplication issues caused by thread-local variables in template classes
   - Removed redundant out-of-class definitions in header files

2. Python Binding Interface Refactoring
   - Introduced BuilderFactory to standardize index builder creation logic
   - Implemented SearcherFactory to support searchers with multiple quantization types
   - Optimized header file reference paths and dependencies
   - Cleaned up legacy implementation code

3. Code Style and Structure Enhancement
   - Removed redundant header inclusions in builder_factory.cpp
   - Cleaned up outdated commented code in sq4_quant.h and sq8_quant.h
   - Standardized header inclusion order in fp32_quant.cpp

Author: Kosthi

python_bindings/bindings.cpp

@@ -7,9 +7,11 @@
 #include <omp.h>
 #endif
 
-#include "hnsw/builder.h"
-#include "hnsw/hnsw.h"
-#include "searcher.h"
+#include "core/interfaces.h"
+#include "graph/builder_factory.h"
+#include "graph/graph.h"
+#include "quantization/fp32_quant.h"
+#include "searcher/searcher.h"
 
 namespace py = pybind11;
 
@@ -62,10 +64,10 @@ void parallel_for(size_t n, int num_threads, Func f) {
 // 3. Graph wrapper
 // -----------------------------------------------------------------------------
 struct Graph {
-  deepsearch::Graph<int> graph;
+  deepsearch::graph::Graph graph;
 
   Graph() = default;
-  explicit Graph(const deepsearch::Graph<int>& graph) : graph(graph) {}
+  explicit Graph(const deepsearch::graph::Graph& graph) : graph(graph) {}
   explicit Graph(const std::string& filename) { graph.load(filename); }
 
   void save(const std::string& filename) { graph.save(filename); }
@@ -76,45 +78,94 @@ struct Graph {
 // 4. Index wrapper
 // -----------------------------------------------------------------------------
 struct Index {
-  std::unique_ptr<deepsearch::Builder> idx;
+  std::unique_ptr<deepsearch::graph::GraphBuilder<float>> builder;
+  deepsearch::core::DistanceType distance_type_;
+  size_t dim_;
 
   Index(const std::string& type, int dim, const std::string& metric, int R = 32,
         int L = 200) {
     if (dim <= 0) throw py::value_error("`dim` must be positive");
     if (R <= 0) throw py::value_error("`R` must be positive");
     if (L < 0) throw py::value_error("`L` must be non-negative");
 
+    dim_ = dim;
+
+    // 解析距离类型
+    if (metric == "L2") {
+      distance_type_ = deepsearch::core::DistanceType::L2;
+    } else if (metric == "IP") {
+      distance_type_ = deepsearch::core::DistanceType::IP;
+    } else if (metric == "COSINE") {
+      distance_type_ = deepsearch::core::DistanceType::COSINE;
+    } else {
+      throw py::value_error("Unknown metric: " + metric);
+    }
+
     if (type == "HNSW") {
-      idx = std::make_unique<deepsearch::Hnsw>(dim, metric, R, L);
+      // 使用BuilderFactory创建HNSW构建器
+      deepsearch::graph::BuilderConfig config;
+      config.M = R;
+      config.ef_construction = L;
+
+      builder = deepsearch::graph::BuilderFactory<float>::create(
+          deepsearch::graph::BuilderType::HNSW, distance_type_, dim_, config);
     } else {
       throw py::value_error("Unknown index type: " + type);
     }
   }
 
   Graph build(py::object data) {
     auto buf = to_buffer<float>(data);
-    if (buf.cols != idx->Dim())
+    if (buf.cols != dim_)
       throw py::value_error("Dimension mismatch: expected " +
-                            std::to_string(idx->Dim()) + ", got " +
+                            std::to_string(dim_) + ", got " +
                             std::to_string(buf.cols));
-    idx->Build(buf.ptr, buf.rows);
-    return Graph(idx->GetGraph());
+
+    auto graph = builder->build(buf.ptr, buf.rows, buf.cols);
+    return Graph(graph);
   }
 };
 
 // -----------------------------------------------------------------------------
 // 5. Searcher wrapper
 // -----------------------------------------------------------------------------
 struct Searcher {
-  std::unique_ptr<deepsearch::SearcherBase> sr;
+  std::unique_ptr<deepsearch::searcher::SearcherBase> searcher;
   ssize_t dim_;
 
   Searcher(const Graph& graph, py::object data, const std::string& metric,
-           int level) {
+           const std::string& quant) {
     auto buf = to_buffer<float>(data);
     dim_ = buf.cols;
-    sr = deepsearch::create_searcher(graph.graph, metric, level);
-    sr->SetData(buf.ptr, buf.rows, buf.cols);
+
+    // 解析距离类型
+    deepsearch::core::DistanceType distance_type;
+    if (metric == "L2") {
+      distance_type = deepsearch::core::DistanceType::L2;
+    } else if (metric == "IP") {
+      distance_type = deepsearch::core::DistanceType::IP;
+    } else if (metric == "COSINE") {
+      distance_type = deepsearch::core::DistanceType::COSINE;
+    } else {
+      throw py::value_error("Unknown metric: " + metric);
+    }
+
+    // 使用SearcherFactory创建FP32搜索器
+    if (quant == "fp32") {
+      searcher = deepsearch::searcher::SearcherFactory::createFP32(
+          graph.graph, distance_type, dim_);
+    } else if (quant == "sq8") {
+      searcher = deepsearch::searcher::SearcherFactory::createSQ8(
+          graph.graph, distance_type, dim_);
+    } else if (quant == "sq4") {
+      searcher = deepsearch::searcher::SearcherFactory::createSQ4(
+          graph.graph, distance_type, dim_);
+    } else {
+      throw py::value_error("Unknown quant: " + quant);
+    }
+
+    // 设置数据
+    searcher->SetData(buf.ptr, buf.rows, buf.cols);
   }
 
   py::array_t<int> search(py::object query, int k) {
@@ -124,7 +175,7 @@ struct Searcher {
                             ")");
 
     int* ids = new int[k];
-    sr->Search(buf.ptr, k, ids);
+    searcher->Search(buf.ptr, k, ids);
 
     py::capsule free_when_done(ids,
                                [](void* f) { delete[] static_cast<int*>(f); });
@@ -145,13 +196,12 @@ struct Searcher {
     int* ids = new int[nq * k];
 
     parallel_for(nq, num_threads, [&](size_t i) {
-      sr->Search(buf.ptr + i * dim_, k, ids + i * k);
+      searcher->Search(buf.ptr + i * dim_, k, ids + i * k);
     });
 
     py::capsule free_when_done(ids,
                                [](void* f) { delete[] static_cast<int*>(f); });
 
-    // 返回二维数组，Python 侧析构时自动调用 capsule
     return py::array_t<int>({(ssize_t)nq, (ssize_t)k},    // shape
                             {(ssize_t)(k * sizeof(int)),  // row stride
                              (ssize_t)(sizeof(int))},     // col stride
@@ -162,12 +212,13 @@ struct Searcher {
 
   void set_ef(int ef) {
     if (ef <= 0) throw py::value_error("`ef` must be positive");
-    sr->SetEf(ef);
+    searcher->SetEf(ef);
   }
 
   void optimize(int num_threads = 0) {
     // Use parallel_for with a single iteration to adjust threads
-    parallel_for(1, num_threads, [&](size_t) { sr->Optimize(num_threads); });
+    parallel_for(1, num_threads,
+                 [&](size_t) { searcher->Optimize(num_threads); });
   }
 };
 
@@ -202,9 +253,10 @@ PYBIND11_MODULE(deepsearch, m) {
            "Build the index from a float array");
 
   py::class_<Searcher>(m, "Searcher")
-      .def(py::init<const Graph&, py::object, const std::string&, int>(),
+      .def(py::init<const Graph&, py::object, const std::string&,
+                    const std::string&>(),
            py::arg("graph"), py::arg("data"), py::arg("metric"),
-           py::arg("level"))
+           py::arg("quant"))
       .def("search", &Searcher::search, py::arg("query"), py::arg("k"),
            "Search a single vector")
       .def("batch_search", &Searcher::batch_search, py::arg("query"),

src/graph/builder_factory.cpp

@@ -3,7 +3,6 @@
 #include <stdexcept>
 
 #include "core/exceptions.h"
-#include "hnsw_builder.cpp"  // 包含模板实现
 
 namespace deepsearch {
 namespace graph {

src/main.cpp

@@ -3,8 +3,8 @@
 #include <iostream>
 #include <string>
 
-#include "builder_factory.h"
-#include "searcher.h"
+#include "graph/builder_factory.h"
+#include "searcher/searcher.h"
 
 using namespace deepsearch::core;
 using namespace deepsearch::graph;

src/quantization/fp32_quant.cpp

@@ -1,12 +1,10 @@
 #include "fp32_quant.h"
 
-#include <core/factory.h>
-#include <distance/computers.h>
-
 #include <cstring>
 #include <stdexcept>
 
 #include "allocator.h"
+#include "distance/computers.h"
 
 namespace deepsearch {
 namespace quantization {
@@ -67,21 +65,22 @@ void FP32Quantizer::prefetch_data(size_t index, int lines) const {
 
 void FP32Quantizer::encode_query(const float* query) {
   ensure_thread_query_initialized();
-  encode(query, query_.get());
+  encode(query, get_query().get());
 }
 
 float FP32Quantizer::compute_query_distance(size_t index) const {
   const float* data_code = reinterpret_cast<const float*>(get_data(index));
-  return distance_computer_->compute(query_.get(), data_code);
+  return distance_computer_->compute(get_query().get(), data_code);
 }
 
 float FP32Quantizer::compute_query_distance(const float* code) const {
-  return distance_computer_->compute(query_.get(), code);
+  return distance_computer_->compute(get_query().get(), code);
 }
 
 void FP32Quantizer::ensure_thread_query_initialized() {
-  if (query_ == nullptr) {
-    query_.reset(static_cast<float*>(alloc64B(d_align * sizeof(float))));
+  auto& query = get_query();
+  if (query == nullptr) {
+    query.reset(static_cast<float*>(alloc64B(d_align * sizeof(float))));
   }
 }
 

src/quantization/quantizer.h

@@ -46,19 +46,19 @@ class QuantizerBase {
     }
   };
 
-  // 将 query_ 改为线程局部存储
-  static thread_local std::unique_ptr<CodeType[], FreeDeleter> query_;
+  // 使用函数返回局部静态变量，完全避免重复符号问题
+  static std::unique_ptr<CodeType[], FreeDeleter>& get_query() {
+    thread_local std::unique_ptr<CodeType[], FreeDeleter> query_ = nullptr;
+    return query_;
+  }
+
+  // query_ 使用线程局部存储，使用 inline 关键字解决重复符号问题
+  // inline static thread_local std::unique_ptr<CodeType[], FreeDeleter> query_;
 
   // 添加线程局部初始化方法
   virtual void ensure_thread_query_initialized() = 0;
 };
 
-// 线程局部变量类外定义
-template <typename InputType, typename CodeType>
-thread_local std::unique_ptr<
-    CodeType[], typename QuantizerBase<InputType, CodeType>::FreeDeleter>
-    QuantizerBase<InputType, CodeType>::query_ = nullptr;
-
 // 距离计算器基类
 template <typename T>
 class DistanceComputerBase {

src/quantization/sq4_quant.cpp

@@ -103,21 +103,22 @@ float SQ4Quantizer::compute_distance(const uint8_t* a, const uint8_t* b) const {
 
 void SQ4Quantizer::encode_query(const float* query) {
   ensure_thread_query_initialized();
-  encode(query, query_.get());
+  encode(query, get_query().get());
 }
 
 float SQ4Quantizer::compute_query_distance(size_t index) const {
   const uint8_t* data_code = reinterpret_cast<const uint8_t*>(get_data(index));
-  return distance_computer_->compute(query_.get(), data_code);
+  return distance_computer_->compute(get_query().get(), data_code);
 }
 
 float SQ4Quantizer::compute_query_distance(const uint8_t* code) const {
-  return distance_computer_->compute(query_.get(), code);
+  return distance_computer_->compute(get_query().get(), code);
 }
 
 void SQ4Quantizer::ensure_thread_query_initialized() {
-  if (query_ == nullptr) {
-    query_.reset(static_cast<uint8_t*>(alloc64B(d_align * sizeof(uint8_t))));
+  auto& query = get_query();
+  if (query == nullptr) {
+    query.reset(static_cast<uint8_t*>(alloc64B(d_align * sizeof(uint8_t))));
   }
 }
 

src/quantization/sq4_quant.h

@@ -3,13 +3,14 @@
 #include <algorithm>
 #include <vector>
 
-#include "fp32_quant.h"
 #include "quantizer.h"
 #include "simd/distance_functions.h"
 
 namespace deepsearch {
 namespace quantization {
 
+class FP32Quantizer;
+
 class SQ4Quantizer : public QuantizerBase<float, uint8_t> {
  public:
   using data_type = uint8_t;
@@ -49,36 +50,37 @@ class SQ4Quantizer : public QuantizerBase<float, uint8_t> {
   // 重排序接口 - 使用FP32量化器进行精排
   template <typename Pool>
   void reorder(const Pool& pool, const float* query, int* dst, int k) const {
-    if (reorder_quantizer_) {
-      // 使用FP32量化器进行精确重排序
-      std::vector<std::pair<int, float>> candidates;
-      candidates.reserve(std::min(k, pool.size()));
-
-      for (int i = 0; i < std::min(k, pool.size()); ++i) {
-        int id = pool.id(i);
-        if (id >= 0) {
-          // 从FP32量化器获取原始浮点数据
-          const float* fp32_data =
-              reinterpret_cast<const float*>(reorder_quantizer_->get_data(id));
-          // 使用FP32量化器的距离计算方法
-          float exact_dist =
-              reorder_quantizer_->compute_distance(query, fp32_data);
-          candidates.emplace_back(id, exact_dist);
-        }
-      }
-
-      std::sort(
-          candidates.begin(), candidates.end(),
-          [](const auto& a, const auto& b) { return a.second < b.second; });
-
-      int result_size = std::min(k, static_cast<int>(candidates.size()));
-      for (int i = 0; i < result_size; ++i) {
-        dst[i] = candidates[i].first;
-      }
-      for (int i = result_size; i < k; ++i) {
-        dst[i] = -1;
-      }
-    }
+    // if (reorder_quantizer_) {
+    //   // 使用FP32量化器进行精确重排序
+    //   std::vector<std::pair<int, float>> candidates;
+    //   candidates.reserve(std::min(k, pool.size()));
+    //
+    //   for (int i = 0; i < std::min(k, pool.size()); ++i) {
+    //     int id = pool.id(i);
+    //     if (id >= 0) {
+    //       // 从FP32量化器获取原始浮点数据
+    //       const float* fp32_data =
+    //           reinterpret_cast<const
+    //           float*>(reorder_quantizer_->get_data(id));
+    //       // 使用FP32量化器的距离计算方法
+    //       float exact_dist =
+    //           reorder_quantizer_->compute_distance(query, fp32_data);
+    //       candidates.emplace_back(id, exact_dist);
+    //     }
+    //   }
+    //
+    //   std::sort(
+    //       candidates.begin(), candidates.end(),
+    //       [](const auto& a, const auto& b) { return a.second < b.second; });
+    //
+    //   int result_size = std::min(k, static_cast<int>(candidates.size()));
+    //   for (int i = 0; i < result_size; ++i) {
+    //     dst[i] = candidates[i].first;
+    //   }
+    //   for (int i = result_size; i < k; ++i) {
+    //     dst[i] = -1;
+    //   }
+    // }
   }
 
  private: