feat: integrate OMEGA adaptive early termination into zvec

.github/workflows/04-android-build.yml

@@ -20,6 +20,8 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v6
+        with:
+          submodules: recursive
 
       - name: Cache dependencies
         uses: actions/cache@v5
@@ -60,7 +62,7 @@ jobs:
       - name: Use host env to compile protoc
         shell: bash
         run: |
-          git submodule update --init
+          git submodule update --init --recursive
           if [ ! -d "build-host" ]; then
             export CCACHE_BASEDIR="$GITHUB_WORKSPACE"
             export CCACHE_NOHASHDIR=1

.github/workflows/scripts/run_vdb.sh

@@ -15,7 +15,7 @@ echo "workspace: $GITHUB_WORKSPACE"
 DB_LABEL_PREFIX="Zvec16c64g-$COMMIT_ID"
 
 # install zvec
-git submodule update --init
+git submodule update --init --recursive
 
 # for debug
 #cd ..
@@ -85,4 +85,4 @@ EOF
         cat prom_metrics.txt
         curl --data-binary @prom_metrics.txt "http://47.93.34.27:9091/metrics/job/benchmarks-${CASE_TYPE}/case_type/${CASE_TYPE}/quantize_type/${QUANTIZE_TYPE}" -v
     done
-done
+done

.gitmodules

@@ -40,6 +40,9 @@
 	path = thirdparty/magic_enum/magic_enum-0.9.7
 	url = https://github.com/Neargye/magic_enum.git
 	ignore = all
+[submodule "thirdparty/omega/OMEGALib"]
+	path = thirdparty/omega/OMEGALib
+	url = https://github.com/driPyf/OMEGALib.git
 [submodule "thirdparty/RaBitQ-Library/RaBitQ-Library-0.1"]
 	path = thirdparty/RaBitQ-Library/RaBitQ-Library-0.1
 	url = https://github.com/VectorDB-NTU/RaBitQ-Library.git

CMakeLists.txt

@@ -65,6 +65,14 @@ message(STATUS "BUILD_C_BINDINGS:${BUILD_C_BINDINGS}")
 option(BUILD_TOOLS "Build tools" ON)
 message(STATUS "BUILD_TOOLS:${BUILD_TOOLS}")
 
+if(ANDROID)
+    option(ZVEC_ENABLE_OMEGA "Build OMEGA support" OFF)
+else()
+    option(ZVEC_ENABLE_OMEGA "Build OMEGA support" ON)
+endif()
+message(STATUS "ZVEC_ENABLE_OMEGA:${ZVEC_ENABLE_OMEGA}")
+add_compile_definitions(ZVEC_ENABLE_OMEGA=$<BOOL:${ZVEC_ENABLE_OMEGA}>)
+
 option(RABITQ_ENABLE_AVX512 "Compile RaBitQ with AVX-512 support" OFF)
 
 if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64|AMD64" AND NOT ANDROID)

cmake/bazel.cmake

@@ -453,7 +453,7 @@ if(CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 7.0)
       "$<$<CXX_COMPILER_ID:Clang>:-Wall;-Wextra;-Wshadow>"
       "$<$<CXX_COMPILER_ID:AppleClang>:-Wall;-Wextra;-Wshadow>"
       "$<$<CXX_COMPILER_ID:GNU>:-Wall;-Wextra;-Wshadow-local;-Wno-misleading-indentation>"
-      "$<$<CXX_COMPILER_ID:MSVC>:/W4>"
+      "$<$<CXX_COMPILER_ID:MSVC>:/W4;/utf-8>"
       ${BAZEL_CC_ASAN_COMPILE_FLAGS}
       ${BAZEL_CC_COVERAGE_COMPILE_FLAGS}
     )
@@ -463,7 +463,7 @@ else()
       "$<$<CXX_COMPILER_ID:Clang>:-Wall;-Wextra;-Wshadow>"
       "$<$<CXX_COMPILER_ID:AppleClang>:-Wall;-Wextra;-Wshadow>"
       "$<$<CXX_COMPILER_ID:GNU>:-Wall;-Wextra;-Wshadow;-Wno-misleading-indentation>"
-      "$<$<CXX_COMPILER_ID:MSVC>:/W4>"
+      "$<$<CXX_COMPILER_ID:MSVC>:/W4;/utf-8>"
       ${BAZEL_CC_ASAN_COMPILE_FLAGS}
       ${BAZEL_CC_COVERAGE_COMPILE_FLAGS}
     )
@@ -484,7 +484,7 @@ set(
     "$<$<CXX_COMPILER_ID:Clang>:-Wall>"
     "$<$<CXX_COMPILER_ID:AppleClang>:-Wall>"
     "$<$<CXX_COMPILER_ID:GNU>:-Wall>"
-    "$<$<CXX_COMPILER_ID:MSVC>:/W3>"
+    "$<$<CXX_COMPILER_ID:MSVC>:/W3;/utf-8>"
     ${BAZEL_CC_ASAN_COMPILE_FLAGS}
     ${BAZEL_CC_COVERAGE_COMPILE_FLAGS}
   )

examples/c++/CMakeLists.txt

@@ -16,6 +16,12 @@ set(ZVEC_INCLUDE_DIR ${CMAKE_BINARY_DIR}/../../../src/include)
 set(ZVEC_LIB_DIR ${CMAKE_BINARY_DIR}/../../../${HOST_BUILD_DIR}/lib)
 set(ZVEC_DEPENDENCY_LIB_DIR ${CMAKE_BINARY_DIR}/../../../${HOST_BUILD_DIR}/external/usr/local/lib)
 
+if(ANDROID)
+    option(ZVEC_ENABLE_OMEGA "Link examples against OMEGA support from the host build" OFF)
+else()
+    option(ZVEC_ENABLE_OMEGA "Link examples against OMEGA support from the host build" ON)
+endif()
+
 # Add include and library search paths
 include_directories(${ZVEC_INCLUDE_DIR})
 link_directories(${ZVEC_LIB_DIR} ${ZVEC_DEPENDENCY_LIB_DIR})
@@ -31,13 +37,48 @@ else()
     set(PROTOBUF_LIB protobuf)
 endif()
 
+if(ZVEC_ENABLE_OMEGA)
+    set(_zvec_omega_search_paths
+        ${ZVEC_DEPENDENCY_LIB_DIR}
+        ${ZVEC_DEPENDENCY_LIB_DIR}/Debug
+        ${ZVEC_DEPENDENCY_LIB_DIR}/Release
+        ${ZVEC_DEPENDENCY_LIB_DIR}/RelWithDebInfo
+        ${ZVEC_DEPENDENCY_LIB_DIR}/MinSizeRel
+    )
+    find_library(ZVEC_OMEGA_LIB
+        NAMES omega
+        PATHS ${_zvec_omega_search_paths}
+        NO_DEFAULT_PATH
+    )
+    find_library(ZVEC_LIGHTGBM_LIB
+        NAMES lib_lightgbm _lightgbm
+        PATHS ${_zvec_omega_search_paths}
+        NO_DEFAULT_PATH
+    )
+    if(NOT ZVEC_OMEGA_LIB OR NOT ZVEC_LIGHTGBM_LIB)
+        message(FATAL_ERROR
+            "ZVEC_ENABLE_OMEGA=ON but failed to locate OMEGA host libraries under "
+            "${ZVEC_DEPENDENCY_LIB_DIR}"
+        )
+    endif()
+endif()
+
 # --- Dependency groups ---
 find_package(Threads REQUIRED)
+find_package(OpenMP QUIET)
 
+set(zvec_openmp_deps)
+if(OpenMP_FOUND AND NOT ANDROID)
+    list(APPEND zvec_openmp_deps OpenMP::OpenMP_CXX)
+endif()
 
 set(zvec_core_deps
     zvec_turbo
+    ${zvec_openmp_deps}
 )
+if(ZVEC_ENABLE_OMEGA)
+    list(APPEND zvec_core_deps ${ZVEC_OMEGA_LIB} ${ZVEC_LIGHTGBM_LIB})
+endif()
 
 if (NOT WIN32)
     set(zvec_ailego_deps
@@ -63,7 +104,11 @@ if (NOT WIN32)
             ${GFLAGS_LIB}
             ${PROTOBUF_LIB}
             lz4
+            ${zvec_openmp_deps}
     )
+    if(ZVEC_ENABLE_OMEGA)
+        list(APPEND zvec_db_deps ${ZVEC_OMEGA_LIB} ${ZVEC_LIGHTGBM_LIB})
+    endif()
 else ()
     # Windows static libraries use different naming conventions
     set(PROTOBUF_LIB libprotobuf)
@@ -90,7 +135,11 @@ else ()
             lz4
             rpcrt4
             shlwapi
+            ${zvec_openmp_deps}
     )
+    if(ZVEC_ENABLE_OMEGA)
+        list(APPEND zvec_db_deps ${ZVEC_OMEGA_LIB} ${ZVEC_LIGHTGBM_LIB})
+    endif()
 endif ()
 
 # --- Create INTERFACE targets for Zvec components ---
@@ -127,6 +176,11 @@ elseif(APPLE)
         zvec-ailego
         ${zvec_core_deps}
     )
+    if(ZVEC_ENABLE_OMEGA)
+        target_link_libraries(zvec-core INTERFACE
+            -Wl,-force_load ${ZVEC_OMEGA_LIB}
+        )
+    endif()
 elseif(ANDROID)
     target_link_libraries(zvec-core INTERFACE
         -Wl,--whole-archive
@@ -161,6 +215,7 @@ if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
     )
 elseif(APPLE)
     target_link_libraries(zvec-db INTERFACE
+        -Wl,-force_load ${ZVEC_LIB_DIR}/libzvec_db.a
         zvec_db
         zvec-core
         zvec-ailego
@@ -208,6 +263,13 @@ target_link_libraries(ailego-example PRIVATE
     zvec-ailego
 )
 
+if(ZVEC_ENABLE_OMEGA)
+    add_executable(omega-example omega/main.cc)
+    target_link_libraries(omega-example PRIVATE
+        zvec-core
+    )
+endif()
+
 # Strip symbols to reduce executable size
 if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID)
     add_custom_command(TARGET db-example POST_BUILD
@@ -219,6 +281,11 @@ if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID)
     add_custom_command(TARGET ailego-example POST_BUILD
         COMMAND ${CMAKE_STRIP} "$<TARGET_FILE:ailego-example>"
         COMMENT "Stripping symbols from ailego-example")
+    if(ZVEC_ENABLE_OMEGA)
+        add_custom_command(TARGET omega-example POST_BUILD
+            COMMAND ${CMAKE_STRIP} "$<TARGET_FILE:omega-example>"
+            COMMENT "Stripping symbols from omega-example")
+    endif()
 endif()
 
 # Optimize for size
@@ -227,4 +294,10 @@ if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID)
                  PROPERTY COMPILE_FLAGS "-Os")
     set_property(TARGET db-example core-example ailego-example
                  PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
+    if(ZVEC_ENABLE_OMEGA)
+        set_property(TARGET omega-example
+                     PROPERTY COMPILE_FLAGS "-Os")
+        set_property(TARGET omega-example
+                     PROPERTY INTERPROCEDURAL_OPTIMIZATION TRUE)
+    endif()
 endif()

examples/c++/omega/main.cc

@@ -0,0 +1,92 @@
+#include <filesystem>
+#include <iostream>
+#include <vector>
+#include <zvec/core/interface/index.h>
+#include <zvec/core/interface/index_factory.h>
+#include <zvec/core/interface/index_param_builders.h>
+
+using namespace zvec::core_interface;
+
+namespace {
+
+constexpr uint32_t kDimension = 32;
+const std::string kIndexPath = "omega_example.index";
+
+BaseIndexParam::Pointer CreateOmegaParam() {
+  auto param = HNSWIndexParamBuilder()
+                   .WithMetricType(MetricType::kInnerProduct)
+                   .WithDataType(DataType::DT_FP32)
+                   .WithDimension(kDimension)
+                   .WithIsSparse(false)
+                   .WithM(8)
+                   .WithEFConstruction(64)
+                   .Build();
+  param->index_type = IndexType::kOMEGA;
+  return param;
+}
+
+}  // namespace
+
+int main() {
+  std::filesystem::remove_all(kIndexPath);
+
+  auto index = IndexFactory::CreateAndInitIndex(*CreateOmegaParam());
+  if (!index) {
+    std::cerr << "failed to create omega index" << std::endl;
+    return 1;
+  }
+
+  if (index->Open(kIndexPath, StorageOptions{StorageOptions::StorageType::kMMAP,
+                                             true}) != 0) {
+    std::cerr << "failed to open omega index" << std::endl;
+    return 1;
+  }
+
+  for (uint32_t doc_id = 0; doc_id < 6; ++doc_id) {
+    std::vector<float> values(kDimension, static_cast<float>(doc_id) / 10.0f);
+    values[0] = 1.0f + static_cast<float>(doc_id);
+    VectorData vector_data;
+    vector_data.vector = DenseVector{values.data()};
+    if (index->Add(vector_data, doc_id) != 0) {
+      std::cerr << "failed to add document " << doc_id << std::endl;
+      return 1;
+    }
+  }
+
+  if (index->Train() != 0) {
+    std::cerr << "failed to train omega index" << std::endl;
+    return 1;
+  }
+
+  std::vector<float> query_values(kDimension, 0.0f);
+  query_values[0] = 1.0f;
+  VectorData query{DenseVector{query_values.data()}};
+
+  auto query_param = OmegaQueryParamBuilder()
+                         .with_topk(3)
+                         .with_fetch_vector(true)
+                         .with_ef_search(32)
+                         .with_target_recall(0.95f)
+                         .build();
+
+  SearchResult result;
+  if (index->Search(query, query_param, &result) != 0) {
+    std::cerr << "failed to search omega index" << std::endl;
+    return 1;
+  }
+
+  std::cout << "omega results: " << result.doc_list_.size() << std::endl;
+  if (result.doc_list_.empty()) {
+    std::cerr << "omega example returned no results" << std::endl;
+    return 1;
+  }
+
+  std::cout << "top result key=" << result.doc_list_[0].key()
+            << " score=" << result.doc_list_[0].score() << std::endl;
+  if (index->Close() != 0) {
+    std::cerr << "failed to close omega index" << std::endl;
+    return 1;
+  }
+
+  return 0;
+}