diff --git a/.github/workflows/04-android-build.yml b/.github/workflows/04-android-build.yml index aa44084c9..89082aead 100644 --- a/.github/workflows/04-android-build.yml +++ b/.github/workflows/04-android-build.yml @@ -23,7 +23,9 @@ jobs: strategy: fail-fast: false matrix: - abi: [x86_64] + # x86_64 runs the emulator suite; arm64-v8a provides compile/link + # coverage for the ABI used by physical Android devices. + abi: [x86_64, arm64-v8a] api: ${{ github.event.inputs.api && fromJSON(format('["{0}"]', github.event.inputs.api)) || fromJSON('["34"]') }} steps: # ── Environment setup ────────────────────────────────────────────── @@ -55,16 +57,22 @@ jobs: uses: android-actions/setup-android@v4 - name: Enable KVM + if: matrix.abi == 'x86_64' run: sudo chmod 666 /dev/kvm || true - - name: Install NDK, emulator and system image + - name: Install NDK and Android platform shell: bash run: | sdkmanager --install \ "ndk;$NDK_VERSION" \ "platform-tools" \ - "platforms;android-${{ matrix.api }}" \ - "emulator" + "platforms;android-${{ matrix.api }}" + + - name: Install emulator and x86_64 system image + if: matrix.abi == 'x86_64' + shell: bash + run: | + sdkmanager --install "emulator" # Install x86_64 system image (try variants in order of availability) sdkmanager --install "system-images;android-${{ matrix.api }};google_apis;x86_64" 2>/dev/null || \ @@ -117,8 +125,12 @@ jobs: -DANDROID_NATIVE_API_LEVEL=${{ matrix.api }} \ -DANDROID_STL=c++_static \ -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_ZVEC_SHARED=OFF \ + -DBUILD_ZVEC_AILEGO_SHARED=OFF \ + -DBUILD_ZVEC_CORE_SHARED=OFF \ -DBUILD_PYTHON_BINDINGS=OFF \ -DBUILD_TOOLS=OFF \ + -DBUILD_CPP_EXAMPLES=ON \ -DENABLE_NATIVE=OFF \ -DAUTO_DETECT_ARCH=OFF \ -DENABLE_WERROR=ON \ @@ -127,7 +139,15 @@ jobs: -DCMAKE_C_COMPILER_LAUNCHER=ccache \ -DCMAKE_CXX_COMPILER_LAUNCHER=ccache - echo "Building all targets..." + if [ "${{ matrix.abi }}" = "arm64-v8a" ]; then + echo "Building focused DiskAnn tests and public C++ examples for arm64-v8a..." + cmake --build "$BUILD_DIR" \ + --target diskann_mobile_compat_test diskann_mobile_collection_test zvec_cpp_examples \ + --parallel + exit 0 + fi + + echo "Building all x86_64 targets..." cmake --build "$BUILD_DIR" --parallel # Discover test targets from ctest metadata @@ -152,6 +172,7 @@ jobs: # ── Step 3: start emulator ───────────────────────────────────────── - name: 'Step 3: Start Android emulator' + if: matrix.abi == 'x86_64' shell: bash run: | AVD_NAME="zvec_test_avd" @@ -246,6 +267,7 @@ jobs: # ── Step 4: run unit tests on emulator ───────────────────────────── - name: 'Step 4: Run unit tests on emulator' + if: matrix.abi == 'x86_64' shell: bash env: BUILD_DIR: build_android_${{ matrix.abi }} @@ -387,47 +409,73 @@ jobs: # ── Step 5: build and run examples ───────────────────────────────── - name: 'Step 5: Build and run examples' + if: matrix.abi == 'x86_64' shell: bash env: BUILD_DIR: build_android_${{ matrix.abi }} run: | - ANDROID_NDK_HOME="$ANDROID_HOME/ndk/$NDK_VERSION" - EXAMPLES_BUILD="examples/c++/build-android-examples-${{ matrix.abi }}" + cmake --build "$BUILD_DIR" --target zvec_cpp_examples --parallel + READELF="$ANDROID_HOME/ndk/$NDK_VERSION/toolchains/llvm/prebuilt/linux-x86_64/bin/llvm-readelf" - cmake -S examples/c++ -B "$EXAMPLES_BUILD" -G Ninja \ - -DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake" \ - -DANDROID_ABI=${{ matrix.abi }} \ - -DANDROID_PLATFORM=android-${{ matrix.api }} \ - -DANDROID_STL=c++_static \ - -DCMAKE_BUILD_TYPE=Release \ - -DHOST_BUILD_DIR="$BUILD_DIR" \ - -DCMAKE_C_COMPILER_LAUNCHER=ccache \ - -DCMAKE_CXX_COMPILER_LAUNCHER=ccache - cmake --build "$EXAMPLES_BUILD" --parallel + for example in ailego-example core-example external-vector-example db-example; do + example_path="$BUILD_DIR/bin/$example" + if [ ! -f "$example_path" ]; then + echo "Missing example binary: $example_path" + exit 1 + fi - # Reuse the shared-library directory from Step 4; push again in - # case Step 4 was skipped or the directory was cleaned. - DEVICE_LIB_DIR="/data/local/tmp/zvec_tests/lib" - adb shell "mkdir -p $DEVICE_LIB_DIR" 2>/dev/null || true - SO_COUNT=0 - while IFS= read -r so_file; do - adb push "$so_file" "$DEVICE_LIB_DIR/$(basename "$so_file")" > /dev/null 2>&1 - SO_COUNT=$((SO_COUNT + 1)) - done < <(find "$BUILD_DIR/lib" -name "*.so" -type f 2>/dev/null) - echo "Pushed $SO_COUNT shared libraries to $DEVICE_LIB_DIR" - - for example in ailego-example core-example db-example; do - if [ -f "$EXAMPLES_BUILD/$example" ]; then - echo "=== Running $example ===" - adb push "$EXAMPLES_BUILD/$example" "/data/local/tmp/$example" > /dev/null 2>&1 - adb shell "chmod 755 /data/local/tmp/$example && cd /data/local/tmp && LD_LIBRARY_PATH=$DEVICE_LIB_DIR ./$example" - adb shell "rm -f /data/local/tmp/$example" + echo "=== Verifying $example is self-contained ===" + dynamic_section=$("$READELF" --dynamic "$example_path") + echo "$dynamic_section" | grep NEEDED || true + if echo "$dynamic_section" | grep -Eq 'libzvec[^]]*\.so|libc\+\+_shared\.so'; then + echo "$example unexpectedly depends on a C++ shared library" + exit 1 + fi + + echo "=== Running $example ===" + adb push "$example_path" "/data/local/tmp/$example" > /dev/null 2>&1 + adb shell "chmod 755 /data/local/tmp/$example && cd /data/local/tmp && ./$example" + adb shell "rm -f /data/local/tmp/$example" + done + + - name: 'Step 3: Verify arm64-v8a artifacts' + if: matrix.abi == 'arm64-v8a' + shell: bash + env: + BUILD_DIR: build_android_${{ matrix.abi }} + run: | + READELF="$ANDROID_HOME/ndk/$NDK_VERSION/toolchains/llvm/prebuilt/linux-x86_64/bin/llvm-readelf" + + for binary in \ + diskann_mobile_compat_test \ + diskann_mobile_collection_test \ + ailego-example \ + core-example \ + external-vector-example \ + db-example; do + binary_path="$BUILD_DIR/bin/$binary" + if [ ! -f "$binary_path" ]; then + echo "Missing arm64-v8a binary: $binary_path" + exit 1 + fi + if ! "$READELF" --file-header "$binary_path" | grep -q 'Machine:.*AArch64'; then + echo "$binary_path is not an AArch64 binary" + exit 1 + fi + done + + for example in ailego-example core-example external-vector-example db-example; do + example_path="$BUILD_DIR/bin/$example" + dynamic_section=$("$READELF" --dynamic "$example_path") + if echo "$dynamic_section" | grep -Eq 'libzvec[^]]*\.so|libc\+\+_shared\.so'; then + echo "$example unexpectedly depends on a C++ shared library" + exit 1 fi done # ── Cleanup ──────────────────────────────────────────────────────── - name: Stop emulator - if: always() + if: matrix.abi == 'x86_64' && always() shell: bash run: | adb emu kill 2>/dev/null || true diff --git a/.github/workflows/06-ios-build.yml b/.github/workflows/06-ios-build.yml index 98a444f62..e31348415 100644 --- a/.github/workflows/06-ios-build.yml +++ b/.github/workflows/06-ios-build.yml @@ -69,8 +69,12 @@ jobs: -DCMAKE_OSX_ARCHITECTURES="${{ matrix.arch }}" \ -DCMAKE_OSX_SYSROOT="$SDK_PATH" \ -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_ZVEC_SHARED=OFF \ + -DBUILD_ZVEC_AILEGO_SHARED=OFF \ + -DBUILD_ZVEC_CORE_SHARED=OFF \ -DBUILD_PYTHON_BINDINGS=OFF \ -DBUILD_TOOLS=OFF \ + -DBUILD_CPP_EXAMPLES=ON \ -DENABLE_WERROR=ON \ -DCMAKE_INSTALL_PREFIX="./install" \ -DGLOBAL_CC_PROTOBUF_PROTOC="$GITHUB_WORKSPACE/build_host/bin/protoc" \ @@ -81,6 +85,20 @@ jobs: cmake --build build_ios_${{ matrix.platform }} --parallel $NPROC + - name: Build public static C++ examples + run: | + NPROC=$(sysctl -n hw.ncpu) + BUILD_DIR=build_ios_${{ matrix.platform }} + cmake --build "$BUILD_DIR" --target zvec_cpp_examples --parallel "$NPROC" + + for example in ailego-example core-example external-vector-example db-example; do + example_binary="$BUILD_DIR/bin/$example.app/$example" + if [ ! -f "$example_binary" ]; then + echo "Missing iOS example binary: $example_binary" + exit 1 + fi + done + - name: Build test targets if: matrix.test_on_simulator run: | diff --git a/CMakeLists.txt b/CMakeLists.txt index 956ee5599..619754a4d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -74,13 +74,22 @@ endif() include_directories(${PROJECT_ROOT_DIR}/src/include) include_directories(${PROJECT_ROOT_DIR}/src) -option(BUILD_ZVEC_SHARED "Build all-in-one C++ shared library libzvec" ON) -option(BUILD_ZVEC_AILEGO_SHARED "Build all-in-one zvec-ailego shared library libzvec_ailego" ON) -option(BUILD_ZVEC_CORE_SHARED "Build all-in-one zvec-core shared library libzvec_core" ON) +set(ZVEC_CPP_SHARED_DEFAULT ON) +if(ANDROID OR IOS) + # A C++ shared library built with a static libc++ cannot safely exchange STL + # objects with a mobile application. Mobile C++ consumers use the static SDK + # targets below; the shared C API remains available through BUILD_C_BINDINGS. + set(ZVEC_CPP_SHARED_DEFAULT OFF) +endif() + +option(BUILD_ZVEC_SHARED "Build all-in-one C++ shared library libzvec" ${ZVEC_CPP_SHARED_DEFAULT}) +option(BUILD_ZVEC_AILEGO_SHARED "Build all-in-one zvec-ailego shared library libzvec_ailego" ${ZVEC_CPP_SHARED_DEFAULT}) +option(BUILD_ZVEC_CORE_SHARED "Build all-in-one zvec-core shared library libzvec_core" ${ZVEC_CPP_SHARED_DEFAULT}) option(BUILD_PYTHON_BINDINGS "Build Python bindings using pybind11" OFF) option(BUILD_C_BINDINGS "Build C bindings" ON) option(BUILD_TOOLS "Build tools" ON) +option(BUILD_CPP_EXAMPLES "Build C++ examples" OFF) message(STATUS "BUILD_ZVEC_SHARED:${BUILD_ZVEC_SHARED}") message(STATUS "BUILD_ZVEC_AILEGO_SHARED:${BUILD_ZVEC_AILEGO_SHARED}") @@ -88,6 +97,7 @@ message(STATUS "BUILD_ZVEC_CORE_SHARED:${BUILD_ZVEC_CORE_SHARED}") message(STATUS "BUILD_PYTHON_BINDINGS:${BUILD_PYTHON_BINDINGS}") message(STATUS "BUILD_C_BINDINGS:${BUILD_C_BINDINGS}") message(STATUS "BUILD_TOOLS:${BUILD_TOOLS}") +message(STATUS "BUILD_CPP_EXAMPLES:${BUILD_CPP_EXAMPLES}") option(RABITQ_ENABLE_AVX512 "Compile RaBitQ with AVX-512 support" OFF) @@ -122,14 +132,16 @@ else() endif() message(STATUS "RABITQ_ARCH_FLAG: ${RABITQ_ARCH_FLAG}") -# DiskAnn support (Linux x86_64 only; libaio loaded at runtime via dlopen) -if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS) +# DiskAnn support. Desktop Linux x86 uses libaio when available, while +# Android and iOS use the portable synchronous pread backend. +if((CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|i686|i386" AND NOT ANDROID AND NOT IOS) + OR ANDROID OR IOS) set(DISKANN_SUPPORTED ON) add_definitions(-DDISKANN_SUPPORTED=1) else() set(DISKANN_SUPPORTED OFF) add_definitions(-DDISKANN_SUPPORTED=0) - message(STATUS "DiskAnn support disabled - only supported on Linux x86_64") + message(STATUS "DiskAnn support disabled - supported on Linux x86 and Android/iOS") endif() message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}") @@ -142,6 +154,10 @@ message(STATUS "USE_OSS_MIRROR:${USE_OSS_MIRROR}") cc_directory(thirdparty) cc_directories(src) +if(BUILD_CPP_EXAMPLES) + add_subdirectory(examples/c++ EXCLUDE_FROM_ALL) +endif() + cc_directories(tests) add_custom_target(clang_tidy_deps DEPENDS zvec_proto ARROW.BUILD glog gflags Lz4.BUILD) diff --git a/cmake/bazel.cmake b/cmake/bazel.cmake index 01d950b79..ee7ad2bfa 100644 --- a/cmake/bazel.cmake +++ b/cmake/bazel.cmake @@ -616,17 +616,24 @@ function(_absolute_paths _RESULT) set(${_RESULT} "${FILEPATHS}" PARENT_SCOPE) endfunction() -## Add both shared and static library +## Add a main library target and an explicit static variant. macro(_add_library _NAME _OPTION) add_library(${_NAME}_objects OBJECT ${_OPTION} ${ARGN}) add_library( ${_NAME}_static STATIC ${_OPTION} $ ) - if(IOS) - # iOS: create the main target as static too (no shared libs on iOS) + if(IOS OR (ANDROID AND ANDROID_STL STREQUAL "c++_static")) add_library( ${_NAME} STATIC ${_OPTION} $ ) + # Keep the two mobile build targets available but give the main archive a + # distinct file name for Ninja. Link-time canonicalization below ensures + # only the main archive is ever whole-archived into an executable. + set_property(TARGET ${_NAME} PROPERTY OUTPUT_NAME ${_NAME}_main) + set_property( + TARGET ${_NAME} PROPERTY ZVEC_CANONICAL_LINK_TARGET ${_NAME}) + set_property( + TARGET ${_NAME}_static PROPERTY ZVEC_CANONICAL_LINK_TARGET ${_NAME}) else() add_library( ${_NAME} SHARED ${_OPTION} $ @@ -671,6 +678,18 @@ endfunction() ## Link libraries function(_target_link_libraries _NAME) + function(_resolve_link_target LIB RESULT_VAR) + if(TARGET ${LIB}) + get_target_property( + CANONICAL_LINK_TARGET ${LIB} ZVEC_CANONICAL_LINK_TARGET) + if(CANONICAL_LINK_TARGET) + set(${RESULT_VAR} ${CANONICAL_LINK_TARGET} PARENT_SCOPE) + return() + endif() + endif() + set(${RESULT_VAR} ${LIB} PARENT_SCOPE) + endfunction() + function(_collect_always_link_libs LIB_LIST RESULT_VAR) if(NOT _COLLECT_ALWAYS_LINK_VISITED) set(_COLLECT_ALWAYS_LINK_VISITED "" PARENT_SCOPE) @@ -678,6 +697,7 @@ function(_target_link_libraries _NAME) set(LOCAL_RESULT "") foreach(LIB ${LIB_LIST}) + _resolve_link_target(${LIB} LIB) if(NOT TARGET ${LIB}) continue() endif() @@ -707,7 +727,10 @@ function(_target_link_libraries _NAME) endif() get_target_property(LINK_LIBS ${LIB} LINK_LIBRARIES) - if(LINK_LIBS) + get_target_property(LIB_TYPE ${LIB} TYPE) + if(LINK_LIBS AND + NOT LIB_TYPE STREQUAL "SHARED_LIBRARY" AND + NOT LIB_TYPE STREQUAL "MODULE_LIBRARY") _collect_always_link_libs("${LINK_LIBS}" LINK_ALWAYS_LINK_LIBS) list(APPEND LOCAL_RESULT ${LINK_ALWAYS_LINK_LIBS}) endif() @@ -717,11 +740,18 @@ function(_target_link_libraries _NAME) set(${RESULT_VAR} "${LOCAL_RESULT}" PARENT_SCOPE) endfunction() - _collect_always_link_libs("${ARGN}" ALL_ALWAYS_LINK_LIBS) + set(INPUT_LIBS "") + foreach(LIB ${ARGN}) + _resolve_link_target(${LIB} RESOLVED_LIB) + list(APPEND INPUT_LIBS ${RESOLVED_LIB}) + endforeach() + list(REMOVE_DUPLICATES INPUT_LIBS) + + _collect_always_link_libs("${INPUT_LIBS}" ALL_ALWAYS_LINK_LIBS) - set(ALL_LIBS_TO_PROCESS ${ARGN}) + set(ALL_LIBS_TO_PROCESS ${INPUT_LIBS}) foreach(ALWAYS_LIB ${ALL_ALWAYS_LINK_LIBS}) - list(FIND ARGN ${ALWAYS_LIB} FOUND_INDEX) + list(FIND INPUT_LIBS ${ALWAYS_LIB} FOUND_INDEX) if(FOUND_INDEX EQUAL -1) list(APPEND ALL_LIBS_TO_PROCESS ${ALWAYS_LIB}) endif() diff --git a/examples/c++/CMakeLists.txt b/examples/c++/CMakeLists.txt index 4e5c703e1..8c7cc2fbd 100644 --- a/examples/c++/CMakeLists.txt +++ b/examples/c++/CMakeLists.txt @@ -1,6 +1,8 @@ cmake_minimum_required(VERSION 3.13) cmake_policy(SET CMP0077 NEW) -project(zvec-example-c++) +if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) + project(zvec-example-c++) +endif() set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) @@ -15,88 +17,110 @@ endif() get_filename_component(ZVEC_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE) set(ZVEC_INCLUDE_DIR ${ZVEC_ROOT_DIR}/src/include) -set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) - include_directories(${ZVEC_INCLUDE_DIR}) -set(ZVEC_LIB_SEARCH_DIRS ${ZVEC_LIB_DIR}) -# Support multi-config builds (MSVC puts libs in Debug/Release subdirectories) -if(CMAKE_BUILD_TYPE) - set(ZVEC_CONFIG_LIB_DIR ${ZVEC_LIB_DIR}/${CMAKE_BUILD_TYPE}) - if(EXISTS "${ZVEC_CONFIG_LIB_DIR}") - list(APPEND ZVEC_LIB_SEARCH_DIRS ${ZVEC_CONFIG_LIB_DIR}) +if(ANDROID OR IOS) + if(NOT TARGET zvec::static OR + NOT TARGET zvec::core_static OR + NOT TARGET zvec::ailego_static) + message(FATAL_ERROR + "Mobile C++ examples must be built from the zvec root with " + "-DBUILD_CPP_EXAMPLES=ON so they use the static SDK targets.") endif() -endif() -if(WIN32) - set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded$<$:Debug>") -endif() -function(zvec_find_shared_library OUT_VAR LIB_NAME) - unset(${OUT_VAR} CACHE) + add_library(zvec-lib INTERFACE) + target_link_libraries(zvec-lib INTERFACE zvec::static) + + add_library(zvec-ailego-lib INTERFACE) + target_link_libraries(zvec-ailego-lib INTERFACE zvec::ailego_static) + + add_library(zvec-core-lib INTERFACE) + target_link_libraries(zvec-core-lib INTERFACE zvec::core_static) +elseif(TARGET zvec_shared AND + TARGET zvec_core_shared AND + TARGET zvec_ailego_shared) + # An in-tree desktop build can link targets directly; the shared-library + # files do not need to exist yet during CMake configuration. + add_library(zvec-lib INTERFACE) + target_link_libraries(zvec-lib INTERFACE zvec_shared) + + add_library(zvec-ailego-lib INTERFACE) + target_link_libraries(zvec-ailego-lib INTERFACE zvec_ailego_shared) + + add_library(zvec-core-lib INTERFACE) + target_link_libraries(zvec-core-lib INTERFACE zvec_core_shared) +else() + set(ZVEC_LIB_DIR ${ZVEC_ROOT_DIR}/${HOST_BUILD_DIR}/lib) + set(ZVEC_LIB_SEARCH_DIRS ${ZVEC_LIB_DIR}) + + # Support multi-config builds (MSVC puts libs in Debug/Release subdirectories) + if(CMAKE_BUILD_TYPE) + set(ZVEC_CONFIG_LIB_DIR ${ZVEC_LIB_DIR}/${CMAKE_BUILD_TYPE}) + if(EXISTS "${ZVEC_CONFIG_LIB_DIR}") + list(APPEND ZVEC_LIB_SEARCH_DIRS ${ZVEC_CONFIG_LIB_DIR}) + endif() + endif() if(WIN32) - find_library(${OUT_VAR} - NAMES ${LIB_NAME}_shared ${LIB_NAME} - PATHS ${ZVEC_LIB_SEARCH_DIRS} - NO_DEFAULT_PATH - NO_CMAKE_FIND_ROOT_PATH - ) - else() - set(ZVEC_ORIGINAL_LIBRARY_SUFFIXES ${CMAKE_FIND_LIBRARY_SUFFIXES}) - if(APPLE) - set(CMAKE_FIND_LIBRARY_SUFFIXES ".dylib") + set(CMAKE_MSVC_RUNTIME_LIBRARY "MultiThreaded$<$:Debug>") + endif() + + function(zvec_find_shared_library OUT_VAR LIB_NAME) + unset(${OUT_VAR} CACHE) + if(WIN32) + find_library(${OUT_VAR} + NAMES ${LIB_NAME}_shared ${LIB_NAME} + PATHS ${ZVEC_LIB_SEARCH_DIRS} + NO_DEFAULT_PATH + NO_CMAKE_FIND_ROOT_PATH + ) else() - set(CMAKE_FIND_LIBRARY_SUFFIXES ".so") + set(ZVEC_ORIGINAL_LIBRARY_SUFFIXES ${CMAKE_FIND_LIBRARY_SUFFIXES}) + if(APPLE) + set(CMAKE_FIND_LIBRARY_SUFFIXES ".dylib") + else() + set(CMAKE_FIND_LIBRARY_SUFFIXES ".so") + endif() + find_library(${OUT_VAR} + NAMES ${LIB_NAME} + PATHS ${ZVEC_LIB_SEARCH_DIRS} + NO_DEFAULT_PATH + NO_CMAKE_FIND_ROOT_PATH + ) + set(CMAKE_FIND_LIBRARY_SUFFIXES "${ZVEC_ORIGINAL_LIBRARY_SUFFIXES}") endif() - find_library(${OUT_VAR} - NAMES ${LIB_NAME} - PATHS ${ZVEC_LIB_SEARCH_DIRS} - NO_DEFAULT_PATH - NO_CMAKE_FIND_ROOT_PATH - ) - set(CMAKE_FIND_LIBRARY_SUFFIXES "${ZVEC_ORIGINAL_LIBRARY_SUFFIXES}") - endif() - set(${OUT_VAR} "${${OUT_VAR}}" PARENT_SCOPE) -endfunction() + set(${OUT_VAR} "${${OUT_VAR}}" PARENT_SCOPE) + endfunction() + + function(zvec_require_shared_library OUT_VAR LIB_NAME) + zvec_find_shared_library(${OUT_VAR} ${LIB_NAME}) + if(NOT ${OUT_VAR}) + message(FATAL_ERROR + "lib${LIB_NAME} shared library was not found in ${ZVEC_LIB_SEARCH_DIRS}. " + "Build zvec first, or pass -DHOST_BUILD_DIR=.") + endif() + set(${OUT_VAR} "${${OUT_VAR}}" PARENT_SCOPE) + endfunction() -function(zvec_require_shared_library OUT_VAR LIB_NAME) - zvec_find_shared_library(${OUT_VAR} ${LIB_NAME}) - if(NOT ${OUT_VAR}) - message(FATAL_ERROR - "lib${LIB_NAME} shared library was not found in ${ZVEC_LIB_SEARCH_DIRS}. " - "Build zvec first, or pass -DHOST_BUILD_DIR=.") - endif() - set(${OUT_VAR} "${${OUT_VAR}}" PARENT_SCOPE) -endfunction() - -zvec_require_shared_library(ZVEC_SHARED_LIBRARY zvec) -zvec_require_shared_library(ZVEC_AILEGO_SHARED_LIBRARY zvec_ailego) -zvec_require_shared_library(ZVEC_CORE_SHARED_LIBRARY zvec_core) - -# --- Create INTERFACE target for libzvec (all-in-one C++ shared library) --- -# libzvec.so/.dylib/.dll already bundles all zvec internal components -# (zvec, zvec_core, zvec_ailego, zvec_turbo), so no individual dependency -# libraries need to be specified by the consumer. -add_library(zvec-lib INTERFACE) -target_link_libraries(zvec-lib INTERFACE "${ZVEC_SHARED_LIBRARY}") - -# --- Create INTERFACE target for libzvec_ailego (ailego-only all-in-one library) --- -# The ailego example intentionally depends only on libzvec_ailego. -add_library(zvec-ailego-lib INTERFACE) -target_link_libraries(zvec-ailego-lib INTERFACE "${ZVEC_AILEGO_SHARED_LIBRARY}") - -# --- Create INTERFACE target for libzvec_core (core-only all-in-one library) --- -# The core example intentionally depends only on libzvec_core. -add_library(zvec-core-lib INTERFACE) -target_link_libraries(zvec-core-lib INTERFACE "${ZVEC_CORE_SHARED_LIBRARY}") + zvec_require_shared_library(ZVEC_SHARED_LIBRARY zvec) + zvec_require_shared_library(ZVEC_AILEGO_SHARED_LIBRARY zvec_ailego) + zvec_require_shared_library(ZVEC_CORE_SHARED_LIBRARY zvec_core) + + # Desktop examples keep using the public all-in-one shared libraries. + add_library(zvec-lib INTERFACE) + target_link_libraries(zvec-lib INTERFACE "${ZVEC_SHARED_LIBRARY}") + + add_library(zvec-ailego-lib INTERFACE) + target_link_libraries(zvec-ailego-lib INTERFACE "${ZVEC_AILEGO_SHARED_LIBRARY}") + + add_library(zvec-core-lib INTERFACE) + target_link_libraries(zvec-core-lib INTERFACE "${ZVEC_CORE_SHARED_LIBRARY}") +endif() # --- Executables --- set(ZVEC_EXAMPLE_TARGETS) add_executable(db-example db/main.cc) target_link_libraries(db-example PRIVATE zvec-lib) -if(ANDROID) - target_link_libraries(db-example PRIVATE log) -endif() list(APPEND ZVEC_EXAMPLE_TARGETS db-example) add_executable(ailego-example ailego/main.cc) @@ -111,6 +135,8 @@ add_executable(external-vector-example core/external_vector_example.cc) target_link_libraries(external-vector-example PRIVATE zvec-core-lib) list(APPEND ZVEC_EXAMPLE_TARGETS external-vector-example) +add_custom_target(zvec_cpp_examples DEPENDS ${ZVEC_EXAMPLE_TARGETS}) + # Strip symbols to reduce executable size if(CMAKE_BUILD_TYPE STREQUAL "Release" AND ANDROID) foreach(ZVEC_EXAMPLE_TARGET ${ZVEC_EXAMPLE_TARGETS}) diff --git a/examples/c/diskann_example.c b/examples/c/diskann_example.c index 5ef80562f..2011de3ff 100644 --- a/examples/c/diskann_example.c +++ b/examples/c/diskann_example.c @@ -21,8 +21,8 @@ * a Vamana graph structure combined with product quantization (PQ) to * achieve high recall with efficient disk I/O. * - * NOTE: DiskANN requires Linux x86_64 with libaio. On other platforms the - * example will compile but the runtime plugin will fail to load. + * NOTE: DiskANN uses libaio when available on Linux x86. Android and iOS use + * the portable synchronous pread backend. * * Workflow demonstrated: * 1. Create collection schema with DiskANN-indexed vector field diff --git a/examples/c/optimized_example.c b/examples/c/optimized_example.c index 28be5c2a2..1acc76eb9 100644 --- a/examples/c/optimized_example.c +++ b/examples/c/optimized_example.c @@ -43,7 +43,7 @@ static float *create_test_vector(size_t dimension) { } for (size_t i = 0; i < dimension; i++) { - vector[i] = (float)rand() / RAND_MAX; + vector[i] = (float)rand() / (float)RAND_MAX; } return vector; @@ -307,4 +307,4 @@ int main() { printf("✓ Optimized example completed\n"); return 0; -} \ No newline at end of file +} diff --git a/scripts/build_android.sh b/scripts/build_android.sh index a1785a4d1..ac8c7a336 100755 --- a/scripts/build_android.sh +++ b/scripts/build_android.sh @@ -89,8 +89,12 @@ cmake -S . -B "$BUILD_DIR" -G Ninja \ -DANDROID_NATIVE_API_LEVEL="$API_LEVEL" \ -DANDROID_STL="c++_static" \ -DCMAKE_BUILD_TYPE="$BUILD_TYPE" \ + -DBUILD_ZVEC_SHARED=OFF \ + -DBUILD_ZVEC_AILEGO_SHARED=OFF \ + -DBUILD_ZVEC_CORE_SHARED=OFF \ -DBUILD_PYTHON_BINDINGS=OFF \ -DBUILD_TOOLS=OFF \ + -DBUILD_CPP_EXAMPLES=ON \ -DENABLE_NATIVE=OFF \ -DAUTO_DETECT_ARCH=OFF \ -DCMAKE_INSTALL_PREFIX="$BUILD_DIR/install" \ @@ -392,4 +396,33 @@ if [ $FAILED -gt 0 ]; then exit 1 fi +echo "" +echo ">>> Step 6: Running statically linked C++ examples..." +cmake --build "$BUILD_DIR" --target zvec_cpp_examples -j"$CORE_COUNT" +READELF=$(find "$ANDROID_NDK_HOME/toolchains/llvm/prebuilt" -type f -name llvm-readelf | head -1) +if [ -z "$READELF" ]; then + echo "ERROR: llvm-readelf was not found in $ANDROID_NDK_HOME" + exit 1 +fi + +for example in ailego-example core-example external-vector-example db-example; do + example_path="$BUILD_DIR/bin/$example" + if [ ! -f "$example_path" ]; then + echo "ERROR: Example binary not found: $example_path" + exit 1 + fi + + dynamic_section=$("$READELF" --dynamic "$example_path") + if echo "$dynamic_section" | grep -Eq 'libzvec[^]]*\.so|libc\+\+_shared\.so'; then + echo "ERROR: $example unexpectedly depends on a C++ shared library" + echo "$dynamic_section" | grep NEEDED || true + exit 1 + fi + + echo " Running $example..." + $ADB_BIN push "$example_path" "/data/local/tmp/$example" > /dev/null 2>&1 + $ADB_BIN shell "chmod 755 /data/local/tmp/$example && cd /data/local/tmp && ./$example" + $ADB_BIN shell "rm -f /data/local/tmp/$example" +done + echo "All tests passed!" diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 2520d9074..526989ce1 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -194,6 +194,89 @@ function(zvec_add_all_in_one_shared TARGET_NAME OUTPUT_NAME) ) endfunction() +# Mobile public C++ API. These build-tree targets keep the application and +# zvec in one C++ runtime when the NDK/iOS toolchain uses a static libc++. +# Whole-archive is required because module registration is performed by static +# initializers that otherwise have no referenced symbol at link time. +function(zvec_add_mobile_static_sdk TARGET_NAME) + cmake_parse_arguments(ZVEC_STATIC_SDK "" "" "LIBS" ${ARGN}) + if(NOT ZVEC_STATIC_SDK_LIBS) + message(FATAL_ERROR "zvec_add_mobile_static_sdk requires LIBS") + endif() + + foreach(ZVEC_STATIC_SDK_LIB ${ZVEC_STATIC_SDK_LIBS}) + if(NOT TARGET ${ZVEC_STATIC_SDK_LIB}) + message(FATAL_ERROR + "Target ${ZVEC_STATIC_SDK_LIB} is required by ${TARGET_NAME}") + endif() + endforeach() + + add_library(${TARGET_NAME} INTERFACE) + target_compile_features(${TARGET_NAME} INTERFACE cxx_std_17) + target_include_directories(${TARGET_NAME} + INTERFACE + $ + $ + ) + + if(IOS) + set(ZVEC_STATIC_SDK_LINK_OPTIONS) + foreach(ZVEC_STATIC_SDK_LIB ${ZVEC_STATIC_SDK_LIBS}) + list(APPEND ZVEC_STATIC_SDK_LINK_OPTIONS + -Wl,-force_load,$ + ) + endforeach() + target_link_options(${TARGET_NAME} + INTERFACE ${ZVEC_STATIC_SDK_LINK_OPTIONS} + ) + target_link_libraries(${TARGET_NAME} + INTERFACE + ${ZVEC_STATIC_SDK_LIBS} + Threads::Threads + ${CMAKE_DL_LIBS} + ) + else() + # Keep whole-archive scoped to the SDK archives themselves. Putting + # these flags in target_link_libraries() also encloses transitive + # dependencies inserted by CMake, which forces both Arrow's bundled + # utf8proc and zvec's standalone utf8proc into the executable. + set(ZVEC_STATIC_SDK_LINK_OPTIONS) + foreach(ZVEC_STATIC_SDK_LIB ${ZVEC_STATIC_SDK_LIBS}) + list(APPEND ZVEC_STATIC_SDK_LINK_OPTIONS + -Wl,--whole-archive,$,--no-whole-archive + ) + endforeach() + target_link_options(${TARGET_NAME} + INTERFACE ${ZVEC_STATIC_SDK_LINK_OPTIONS} + ) + target_link_libraries(${TARGET_NAME} + INTERFACE + ${ZVEC_STATIC_SDK_LIBS} + Threads::Threads + ${CMAKE_DL_LIBS} + ) + if(ANDROID) + target_link_libraries(${TARGET_NAME} INTERFACE log) + endif() + endif() +endfunction() + +if(ANDROID OR IOS) + zvec_add_mobile_static_sdk(zvec_static_sdk + LIBS zvec zvec_core zvec_ailego zvec_turbo + ) + zvec_add_mobile_static_sdk(zvec_core_static_sdk + LIBS zvec_core zvec_ailego zvec_turbo + ) + zvec_add_mobile_static_sdk(zvec_ailego_static_sdk + LIBS zvec_ailego + ) + + add_library(zvec::static ALIAS zvec_static_sdk) + add_library(zvec::core_static ALIAS zvec_core_static_sdk) + add_library(zvec::ailego_static ALIAS zvec_ailego_static_sdk) +endif() + if(BUILD_ZVEC_AILEGO_SHARED) zvec_add_all_in_one_shared(zvec_ailego_shared zvec_ailego LIBS diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index f874eba62..faf02e1ba 100644 --- a/src/core/algorithm/CMakeLists.txt +++ b/src/core/algorithm/CMakeLists.txt @@ -17,7 +17,7 @@ else() # Empty stub library for unsupported platforms file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/diskann_stub.cc "// Stub implementation for unsupported platforms\n" - "// DiskAnn only supports Linux x86_64\n" + "// DiskAnn is unavailable on this platform\n" "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" ) diff --git a/src/core/algorithm/diskann/diskann_context.cc b/src/core/algorithm/diskann/diskann_context.cc index f13affb74..3779b9011 100644 --- a/src/core/algorithm/diskann/diskann_context.cc +++ b/src/core/algorithm/diskann/diskann_context.cc @@ -24,7 +24,9 @@ namespace core { DiskAnnContext::DiskAnnContext(const IndexMeta &meta, const IndexMetric::Pointer &measure, const DiskAnnEntity::Pointer &entity) - : dc_(entity.get(), measure, meta.dimension()), entity_{entity} {} + : IndexContext(measure), + dc_(entity.get(), measure, meta.dimension()), + entity_{entity} {} int DiskAnnContext::init(ContextType type, uint32_t graph_degree, uint32_t pq_chunk_num, uint32_t element_size) { @@ -139,4 +141,4 @@ int DiskAnnContext::update_context(ContextType type, const IndexMeta &meta, } } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index ce8c0ca0c..0b845f34f 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -273,6 +273,21 @@ class DiskAnnContext : public IndexContext, return group_num_ > 0; } + //! Preserve query options when a pooled DiskAnn context is recreated for a + //! different index whose buffers have a different layout. + void copy_query_state_from(const DiskAnnContext &other) { + IndexContext::copy_query_state_from(other); + topk_ = other.topk_; + list_size_ = other.list_size_; + group_topk_ = other.group_topk_; + group_num_ = other.group_num_; + fetch_vector_ = other.fetch_vector_; + debug_mode_ = other.debug_mode_; + topk_heap_.clear(); + topk_heap_.limit(topk_); + group_topk_heaps_.clear(); + } + //! Set group params void set_group_params(uint32_t group_num, uint32_t group_topk) override { group_num_ = group_num; diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index cde0755e8..2932e6c6e 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -28,7 +28,7 @@ namespace zvec { namespace core { -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) typedef struct io_event io_event_t; typedef struct iocb iocb_t; @@ -39,7 +39,7 @@ static std::once_flag g_io_backend_log_once; #endif void log_diskann_io_backend() { -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) auto &backend = ailego::IOBackend::Instance(); if (backend.is_pread()) { LOG_WARN( @@ -55,7 +55,7 @@ void log_diskann_io_backend() { } int setup_io_ctx(IOContext &ctx) { -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) std::call_once(g_io_backend_log_once, log_diskann_io_backend); if (ailego::IOBackend::Instance().is_pread()) { return 0; @@ -64,12 +64,13 @@ int setup_io_ctx(IOContext &ctx) { return ret; #else + (void)ctx; return 0; #endif } int destroy_io_ctx(IOContext &ctx) { -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) if (ailego::IOBackend::Instance().is_pread() || ctx == nullptr) { return 0; } @@ -80,6 +81,7 @@ int destroy_io_ctx(IOContext &ctx) { return ret; #else + (void)ctx; return 0; #endif } @@ -102,7 +104,7 @@ static int execute_io_pread(int fd, std::vector &read_reqs) { return 0; } -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) // io_getevents() should only fail permanently for an invalid context or // invalid arguments. If that happens after submission, io_destroy() is the // only safe way to quiesce the context before synchronous I/O touches the same @@ -256,12 +258,14 @@ int execute_io_libaio(IOContext &ctx, int fd, int execute_io(IOContext &ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) { -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) if (ailego::IOBackend::Instance().is_pread() || ctx == nullptr) { return execute_io_pread(fd, read_reqs); } return execute_io_libaio(ctx, fd, read_reqs, n_retries); #else + (void)ctx; + (void)n_retries; return execute_io_pread(fd, read_reqs); #endif } @@ -294,7 +298,7 @@ IOContext &LinuxAlignedFileReader::get_ctx() { } void LinuxAlignedFileReader::register_thread() { -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) auto thread_id = std::this_thread::get_id(); std::unique_lock lk(ctx_mut); if (ctx_map.find(thread_id) != ctx_map.end()) { @@ -330,7 +334,7 @@ void LinuxAlignedFileReader::register_thread() { } void LinuxAlignedFileReader::deregister_thread() { -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) auto thread_id = std::this_thread::get_id(); IOContext ctx; @@ -355,7 +359,7 @@ void LinuxAlignedFileReader::deregister_thread() { } void LinuxAlignedFileReader::deregister_all_threads() { -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) std::unique_lock lk(ctx_mut); bool aio_available = ailego::IOBackend::Instance().available() != ailego::IOBackendType::kPread; @@ -372,13 +376,13 @@ void LinuxAlignedFileReader::deregister_all_threads() { void LinuxAlignedFileReader::open(const std::string &fname) { int flags = O_RDONLY; -#if defined(__linux__) || defined(__linux) +#if defined(__linux__) && !defined(__ANDROID__) flags |= O_DIRECT | O_LARGEFILE; #endif this->file_desc = ::open(fname.c_str(), flags); -#if defined(__linux__) || defined(__linux) +#if defined(__linux__) && !defined(__ANDROID__) // O_DIRECT may not be supported on all filesystems (e.g. tmpfs, overlay). // Fall back to regular buffered I/O when it fails. if (this->file_desc == -1) { diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index a1cb7c91a..6693eb33c 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -17,7 +17,7 @@ #include -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) #include // dlopen-based libaio wrapper #endif @@ -30,7 +30,7 @@ namespace zvec { namespace core { -#if (defined(__linux) || defined(__linux__)) +#if defined(__linux__) && !defined(__ANDROID__) typedef io_context_t IOContext; #else typedef uint32_t IOContext; @@ -40,7 +40,8 @@ int setup_io_ctx(IOContext &ctx); int destroy_io_ctx(IOContext &ctx); // Log the current DiskAnn I/O backend status (async vs. synchronous pread). -// Probes the backend on first call. No-op on non-Linux platforms. +// Probes the backend on first call. No-op outside desktop Linux; Android and +// iOS always use synchronous pread. void log_diskann_io_backend(); struct AlignedRead { diff --git a/src/core/algorithm/diskann/diskann_pq_trainer.cc b/src/core/algorithm/diskann/diskann_pq_trainer.cc index 73ca01656..27ce1f5e6 100644 --- a/src/core/algorithm/diskann/diskann_pq_trainer.cc +++ b/src/core/algorithm/diskann/diskann_pq_trainer.cc @@ -149,19 +149,26 @@ int DiskAnnPqTrainer::convert_pivot_data( for (size_t cluster = 0; cluster < num_centers; ++cluster) { size_t idx = chunk * num_centers + cluster; - T *pivot_data_ptr = reinterpret_cast(&(full_pivot_data[0])) + - cluster * dim + chunk_offsets[chunk]; - const T *feature_ptr = - reinterpret_cast(centroids[idx].feature()); - for (size_t d = 0; d <= chunk_dims[chunk]; ++d) { - pivot_data_ptr[d] = feature_ptr[d]; - } + uint8_t *pivot_data_ptr = + full_pivot_data.data() + + (cluster * dim + chunk_offsets[chunk]) * sizeof(T); + std::memcpy(pivot_data_ptr, centroids[idx].feature(), + chunk_dims[chunk] * sizeof(T)); } } return 0; } +template int DiskAnnPqTrainer::convert_pivot_data( + const IndexMeta &, uint32_t, uint32_t, const std::vector &, + const std::vector &, IndexCluster::CentroidList &, + std::vector &); +template int DiskAnnPqTrainer::convert_pivot_data( + const IndexMeta &, uint32_t, uint32_t, const std::vector &, + const std::vector &, IndexCluster::CentroidList &, + std::vector &); + int DiskAnnPqTrainer::train_pq(IndexThreads::Pointer threads, const IndexMeta &meta, std::string &train_data, size_t num_train, uint32_t num_centers, diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc index a34c546e5..1b80043fe 100644 --- a/src/core/algorithm/diskann/diskann_searcher.cc +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "diskann_searcher.h" +#include #include "diskann_context.h" #include "diskann_indexer.h" #include "diskann_params.h" @@ -145,14 +146,15 @@ int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, // with different element sizes (e.g., fp16 vs fp32), the cached context has // undersized buffers. Recreate it to ensure correct buffer allocations. if (ctx->magic() != magic_) { - uint32_t saved_topk = ctx->topk(); + auto previous_context = std::move(context); + auto *previous_ctx = dynamic_cast(previous_context.get()); context = create_context(); if (!context) { LOG_ERROR("Failed to recreate context for current streamer"); return IndexError_Runtime; } ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + ctx->copy_query_state_from(*previous_ctx); } ctx->clear(); @@ -192,14 +194,15 @@ int DiskAnnSearcher::search_bf_impl(const void *query, if (ctx->magic() != magic_) { //! context is created by another searcher or streamer, recreate it //! to ensure buffers are correctly sized for this index's parameters. - uint32_t saved_topk = ctx->topk(); + auto previous_context = std::move(context); + auto *previous_ctx = dynamic_cast(previous_context.get()); context = create_context(); if (!context) { LOG_ERROR("Failed to recreate context for current streamer"); return IndexError_Runtime; } ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + ctx->copy_query_state_from(*previous_ctx); } ctx->clear(); @@ -245,14 +248,15 @@ int DiskAnnSearcher::search_bf_by_p_keys_impl( if (ctx->magic() != magic_) { //! context is created by another searcher or streamer, recreate it //! to ensure buffers are correctly sized for this index's parameters. - uint32_t saved_topk = ctx->topk(); + auto previous_context = std::move(context); + auto *previous_ctx = dynamic_cast(previous_context.get()); context = create_context(); if (!context) { LOG_ERROR("Failed to recreate context for current streamer"); return IndexError_Runtime; } ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + ctx->copy_query_state_from(*previous_ctx); } ctx->clear(); diff --git a/src/core/algorithm/diskann/diskann_searcher_entity.cc b/src/core/algorithm/diskann/diskann_searcher_entity.cc index c9e49deba..9c9d1128f 100644 --- a/src/core/algorithm/diskann/diskann_searcher_entity.cc +++ b/src/core/algorithm/diskann/diskann_searcher_entity.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "diskann_searcher_entity.h" +#include namespace zvec { namespace core { @@ -398,7 +399,7 @@ const void *DiskAnnSearcherEntity::get_vector(diskann_id_t id) const { const void *vec; if (ailego_unlikely(vector_segment_->read(total_offset, &vec, read_size) != read_size)) { - LOG_ERROR("Read vector from segment failed, id: %u, offset: %lu", id, + LOG_ERROR("Read vector from segment failed, id: %u, offset: %" PRIu64, id, total_offset); return nullptr; } diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 82e97dcd6..4c3680eee 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -13,6 +13,7 @@ // limitations under the License. #include "diskann_streamer.h" +#include #include "diskann_context.h" #include "diskann_index_provider.h" #include "diskann_indexer.h" @@ -144,14 +145,15 @@ int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, // with different element sizes (e.g., fp16 vs fp32), the cached context has // undersized buffers. Recreate it to ensure correct buffer allocations. if (ctx->magic() != magic_) { - uint32_t saved_topk = ctx->topk(); + auto previous_context = std::move(context); + auto *previous_ctx = dynamic_cast(previous_context.get()); context = create_context(); if (!context) { LOG_ERROR("Failed to recreate context for current streamer"); return IndexError_Runtime; } ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + ctx->copy_query_state_from(*previous_ctx); } ctx->clear(); @@ -187,14 +189,15 @@ int DiskAnnStreamer::search_bf_impl(const void *query, if (ctx->magic() != magic_) { //! context is created by another searcher or streamer, recreate it //! to ensure buffers are correctly sized for this index's parameters. - uint32_t saved_topk = ctx->topk(); + auto previous_context = std::move(context); + auto *previous_ctx = dynamic_cast(previous_context.get()); context = create_context(); if (!context) { LOG_ERROR("Failed to recreate context for current streamer"); return IndexError_Runtime; } ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + ctx->copy_query_state_from(*previous_ctx); } ctx->clear(); @@ -240,14 +243,15 @@ int DiskAnnStreamer::search_bf_by_p_keys_impl( if (ctx->magic() != magic_) { //! context is created by another searcher or streamer, recreate it //! to ensure buffers are correctly sized for this index's parameters. - uint32_t saved_topk = ctx->topk(); + auto previous_context = std::move(context); + auto *previous_ctx = dynamic_cast(previous_context.get()); context = create_context(); if (!context) { LOG_ERROR("Failed to recreate context for current streamer"); return IndexError_Runtime; } ctx = dynamic_cast(context.get()); - ctx->set_topk(saved_topk); + ctx->copy_query_state_from(*previous_ctx); } ctx->clear(); @@ -338,6 +342,7 @@ IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { } ctx->set_list_size(list_size_); + ctx->set_magic(magic_); return Context::Pointer(ctx); } diff --git a/src/core/algorithm/diskann/diskann_util.h b/src/core/algorithm/diskann/diskann_util.h index a02130bf0..cb6350042 100644 --- a/src/core/algorithm/diskann/diskann_util.h +++ b/src/core/algorithm/diskann/diskann_util.h @@ -13,6 +13,7 @@ // limitations under the License. #pragma once +#include #include #include #include "diskann_entity.h" @@ -35,7 +36,13 @@ class DiskAnnUtil { } static inline void alloc_aligned(void **ptr, size_t size, size_t align) { - *ptr = ::aligned_alloc(align, size); + if (ptr == nullptr) { + return; + } + *ptr = nullptr; + if (size == 0 || ::posix_memalign(ptr, align, size) != 0) { + *ptr = nullptr; + } } static inline void free_aligned(void *ptr) { @@ -218,4 +225,4 @@ class NeighborPriorityQueue { }; } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index e377f5342..b21ae76b0 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -27,7 +27,7 @@ namespace zvec::core_interface { int DiskAnnIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { (void)param; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } @@ -35,24 +35,24 @@ int DiskAnnIndex::Open(const std::string &file_path, StorageOptions storage_options) { (void)file_path; (void)storage_options; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } int DiskAnnIndex::GenerateHolder() { - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } int DiskAnnIndex::Add(const VectorData &vector, uint32_t doc_id) { (void)vector; (void)doc_id; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } int DiskAnnIndex::Train() { - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } @@ -60,7 +60,7 @@ int DiskAnnIndex::_dense_fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { (void)doc_id; (void)vector_data_buffer; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } @@ -70,7 +70,7 @@ int DiskAnnIndex::_prepare_for_search( (void)query; (void)search_param; (void)context; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } @@ -80,7 +80,7 @@ int DiskAnnIndex::Merge(const std::vector &indexes, (void)indexes; (void)filter; (void)options; - LOG_ERROR("DiskAnn is not supported on this platform (Linux x86_64 only)"); + LOG_ERROR("DiskAnn is not supported on this platform"); return core::IndexError_Unsupported; } @@ -291,6 +291,13 @@ int DiskAnnIndex::_prepare_for_search( } context->set_topk(diskann_search_param->topk); + context->set_fetch_vector(diskann_search_param->fetch_vector); + if (diskann_search_param->filter) { + context->set_filter(std::move(*diskann_search_param->filter)); + } + if (diskann_search_param->radius > 0.0f) { + context->set_threshold(diskann_search_param->radius); + } // Propagate the query-time beam-search list size into the context. Must be // at least topk to keep enough candidates for a correct result. diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index 424f98f71..c4b171841 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -198,7 +198,8 @@ Status FieldSchema::validate() const { } if (index_params_->type() == IndexType::DISKANN) { - // DiskAnn requires Linux x86_64/i686/i386. The CMake variable + // DiskAnn supports Linux x86_64/i686/i386 and Android/iOS. The CMake + // variable // DISKANN_SUPPORTED (defined in the top-level CMakeLists.txt) is the // single source of truth for platform eligibility — it is also used by // index_factory.cc to conditionally compile the DiskAnn index @@ -210,7 +211,8 @@ Status FieldSchema::validate() const { // back to synchronous pread() with degraded performance. #if !DISKANN_SUPPORTED return Status::NotSupported( - "DiskAnn is not supported on this platform (Linux x86_64 only)"); + "DiskAnn is not supported on this platform (supported on Linux " + "x86 and Android/iOS)"); #endif } diff --git a/src/include/zvec/core/framework/index_context.h b/src/include/zvec/core/framework/index_context.h index 141421453..f153106e3 100644 --- a/src/include/zvec/core/framework/index_context.h +++ b/src/include/zvec/core/framework/index_context.h @@ -258,6 +258,24 @@ class IndexContext { return profiler_; } + protected: + //! Copy query-scoped state when a pooled context must be recreated for a + //! different index instance. Derived contexts remain responsible for their + //! own query parameters. + void copy_query_state_from(const IndexContext &other) { + filter_ = other.filter_; + group_by_ = other.group_by_; + threshold_ = other.threshold_; + if (threshold_ != std::numeric_limits::max()) { + if (other.index_metric_ && other.index_metric_->support_normalize()) { + other.index_metric_->normalize(&threshold_); + } + if (index_metric_ && index_metric_->support_normalize()) { + index_metric_->denormalize(&threshold_); + } + } + } + private: //! Members IndexFilter filter_{}; diff --git a/tests/c/CMakeLists.txt b/tests/c/CMakeLists.txt index 9f40ef9ca..f2c3ad850 100644 --- a/tests/c/CMakeLists.txt +++ b/tests/c/CMakeLists.txt @@ -18,7 +18,7 @@ file(GLOB_RECURSE ALL_TEST_SRCS *_test.c) foreach(CC_SRCS ${ALL_TEST_SRCS}) get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) - cc_gtest( + cc_test( NAME ${CC_TARGET} STRICT LIBS zvec_c_api diff --git a/tests/c/c_api_test.c b/tests/c/c_api_test.c index 7365dcf1a..8a9d7cbc9 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -474,11 +474,11 @@ void test_schema_edge_cases(void) { // Test 4: NULL schema parameter handling for all functions zvec_error_code_t err; const char **test_names = NULL; - size_t test_count = 0; + size_t field_name_count = 0; err = zvec_collection_schema_get_all_field_names(NULL, &test_names, - &test_count); + &field_name_count); TEST_ASSERT(err == ZVEC_ERROR_INVALID_ARGUMENT); - TEST_ASSERT(test_count == 0); + TEST_ASSERT(field_name_count == 0); const zvec_field_schema_t *null_field = zvec_collection_schema_get_field(NULL, "test"); @@ -5109,7 +5109,7 @@ void test_performance_benchmarks(void) { // Create random vector float vec[128]; for (int j = 0; j < 128; j++) { - vec[j] = (float)rand() / RAND_MAX; + vec[j] = (float)rand() / (float)RAND_MAX; } zvec_doc_add_field_by_value(batch_docs[i], "vec", ZVEC_DATA_TYPE_VECTOR_FP32, vec, @@ -5150,7 +5150,7 @@ void test_performance_benchmarks(void) { // Test query performance float query_vec[128]; for (int i = 0; i < 128; i++) { - query_vec[i] = (float)rand() / RAND_MAX; + query_vec[i] = (float)rand() / (float)RAND_MAX; } zvec_vector_query_t *query = zvec_vector_query_create(); diff --git a/tests/c/utils.c b/tests/c/utils.c index 61c118849..dfa651d28 100644 --- a/tests/c/utils.c +++ b/tests/c/utils.c @@ -725,12 +725,12 @@ zvec_doc_t *zvec_test_create_doc_null(uint64_t doc_id, break; } - if (err != ZVEC_OK) { // Free field names array before returning if (field_names) { - for (size_t i = 0; i < field_count; i++) { - free((char *)field_names[i]); + for (size_t cleanup_index = 0; cleanup_index < field_count; + cleanup_index++) { + free((char *)field_names[cleanup_index]); } free(field_names); } diff --git a/tests/core/algorithm/diskann/CMakeLists.txt b/tests/core/algorithm/diskann/CMakeLists.txt index e6ad1af12..141112c86 100644 --- a/tests/core/algorithm/diskann/CMakeLists.txt +++ b/tests/core/algorithm/diskann/CMakeLists.txt @@ -2,6 +2,14 @@ include(${PROJECT_ROOT_DIR}/cmake/bazel.cmake) file(GLOB_RECURSE ALL_TEST_SRCS *_test.cc) +# The full DiskAnn suite repeatedly builds 10k-vector indexes and is intended +# for desktop CI. Mobile CI runs a focused compatibility test that covers the +# portable I/O path, failure recovery, concurrency, and an end-to-end +# build/dump/load/search cycle. +if(ANDROID OR IOS) + list(FILTER ALL_TEST_SRCS INCLUDE REGEX "diskann_mobile_compat_test\\.cc$") +endif() + foreach(CC_SRCS ${ALL_TEST_SRCS}) get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) cc_gtest( @@ -11,4 +19,4 @@ foreach(CC_SRCS ${ALL_TEST_SRCS}) SRCS ${CC_SRCS} INCS . ${PROJECT_ROOT_DIR}/src/core ${PROJECT_ROOT_DIR}/src/core/algorithm/diskann ) -endforeach() \ No newline at end of file +endforeach() diff --git a/tests/core/algorithm/diskann/diskann_mobile_compat_test.cc b/tests/core/algorithm/diskann/diskann_mobile_compat_test.cc new file mode 100644 index 000000000..66e6828d7 --- /dev/null +++ b/tests/core/algorithm/diskann/diskann_mobile_compat_test.cc @@ -0,0 +1,416 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "diskann_builder.h" +#include "diskann_file_reader.h" +#include "diskann_pq_trainer.h" +#include "diskann_util.h" + +namespace zvec::core { +namespace { + +class TemporaryFile { + public: + TemporaryFile() : fd_(::mkstemp(path_)) {} + + ~TemporaryFile() { + if (fd_ >= 0) { + ::close(fd_); + } + ::unlink(path_); + } + + TemporaryFile(const TemporaryFile &) = delete; + TemporaryFile &operator=(const TemporaryFile &) = delete; + + int fd() const { + return fd_; + } + + const char *path() const { + return path_; + } + + void release_descriptor_and_unlink() { + if (fd_ >= 0) { + ::close(fd_); + fd_ = -1; + } + ::unlink(path_); + } + + private: + char path_[64] = "DiskAnnMobileCompatTest.XXXXXX"; + int fd_{-1}; +}; + +TEST(DiskAnnMobileCompatTest, AlignedAllocationSupportsUnroundedSize) { + constexpr size_t kSize = 400; + constexpr size_t kAlignment = 256; + + void *buffer = nullptr; + DiskAnnUtil::alloc_aligned(&buffer, kSize, kAlignment); + + ASSERT_NE(buffer, nullptr); + EXPECT_EQ(reinterpret_cast(buffer) % kAlignment, 0u); + std::memset(buffer, 0xa5, kSize); + DiskAnnUtil::free_aligned(buffer); +} + +template +void ExpectExactPqPivotCopy(IndexMeta::DataType data_type) { + constexpr uint32_t kDimension = 4; + constexpr uint32_t kCenterCount = 2; + constexpr uint32_t kChunkCount = 2; + const std::vector chunk_dims{2, 2}; + const std::vector chunk_offsets{0, 2, 4}; + const std::array, 4> values{{ + {{1.0F, 2.0F}}, + {{5.0F, 6.0F}}, + {{3.0F, 4.0F}}, + {{7.0F, 8.0F}}, + }}; + + IndexCluster::CentroidList centroids(values.size()); + for (size_t i = 0; i < values.size(); ++i) { + const std::array feature{{T(values[i][0]), T(values[i][1])}}; + centroids[i].set_feature(feature.data(), sizeof(feature)); + } + + IndexMeta meta(data_type, kDimension); + std::vector pivots; + ASSERT_EQ(DiskAnnPqTrainer::convert_pivot_data( + meta, kCenterCount, kChunkCount, chunk_dims, chunk_offsets, + centroids, pivots), + 0); + ASSERT_EQ(pivots.size(), kCenterCount * meta.element_size()); + + std::array actual{}; + std::memcpy(actual.data(), pivots.data(), pivots.size()); + for (size_t i = 0; i < actual.size(); ++i) { + EXPECT_FLOAT_EQ(static_cast(actual[i]), static_cast(i + 1)); + } +} + +TEST(DiskAnnMobileCompatTest, PqPivotConversionCopiesExactChunkWidths) { + ExpectExactPqPivotCopy(IndexMeta::DataType::DT_FP32); + ExpectExactPqPivotCopy(IndexMeta::DataType::DT_FP16); +} + +TEST(DiskAnnMobileCompatTest, PortableReaderReadsAlignedBatch) { + constexpr size_t kBlockSize = 4096; + constexpr size_t kBlockCount = 2; + constexpr size_t kDataSize = kBlockSize * kBlockCount; + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + std::vector expected(kDataSize); + std::fill(expected.begin(), expected.begin() + kBlockSize, 0x3c); + std::fill(expected.begin() + kBlockSize, expected.end(), 0xc3); + ASSERT_EQ(::pwrite(file.fd(), expected.data(), expected.size(), 0), + static_cast(expected.size())); + + void *output = nullptr; + DiskAnnUtil::alloc_aligned(&output, kDataSize, kBlockSize); + ASSERT_NE(output, nullptr); + std::memset(output, 0, kDataSize); + + LinuxAlignedFileReader reader; + reader.open(file.path()); + IOContext context{}; + std::vector requests; + requests.emplace_back(0, kBlockSize, output); + requests.emplace_back(kBlockSize, kBlockSize, + static_cast(output) + kBlockSize); + + EXPECT_EQ(reader.read(requests, context), 0); + EXPECT_EQ(std::memcmp(output, expected.data(), expected.size()), 0); + + reader.close(); + DiskAnnUtil::free_aligned(output); +} + +TEST(DiskAnnMobileCompatTest, PortableReaderRejectsShortRead) { + constexpr size_t kBlockSize = 4096; + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + std::vector expected(kBlockSize, 0x5a); + ASSERT_EQ(::pwrite(file.fd(), expected.data(), expected.size(), 0), + static_cast(expected.size())); + + void *output = nullptr; + DiskAnnUtil::alloc_aligned(&output, kBlockSize * 2, kBlockSize); + ASSERT_NE(output, nullptr); + + LinuxAlignedFileReader reader; + reader.open(file.path()); + IOContext context{}; + std::vector requests; + requests.emplace_back(0, kBlockSize * 2, output); + + EXPECT_NE(reader.read(requests, context), 0); + + requests.clear(); + requests.emplace_back(0, kBlockSize, output); + EXPECT_EQ(reader.read(requests, context), 0); + EXPECT_EQ(std::memcmp(output, expected.data(), expected.size()), 0); + + reader.close(); + DiskAnnUtil::free_aligned(output); +} + +TEST(DiskAnnMobileCompatTest, PortableReaderRecoversAfterOpenFailure) { + constexpr size_t kBlockSize = 4096; + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + std::vector expected(kBlockSize, 0x6b); + ASSERT_EQ(::pwrite(file.fd(), expected.data(), expected.size(), 0), + static_cast(expected.size())); + + void *output = nullptr; + DiskAnnUtil::alloc_aligned(&output, kBlockSize, kBlockSize); + ASSERT_NE(output, nullptr); + + LinuxAlignedFileReader reader; + reader.open("DiskAnnMobileCompatTest.missing"); + IOContext context{}; + std::vector requests; + requests.emplace_back(0, kBlockSize, output); + EXPECT_NE(reader.read(requests, context), 0); + + reader.open(file.path()); + EXPECT_EQ(reader.read(requests, context), 0); + EXPECT_EQ(std::memcmp(output, expected.data(), expected.size()), 0); + + reader.close(); + DiskAnnUtil::free_aligned(output); +} + +TEST(DiskAnnMobileCompatTest, PortableReaderSupportsConcurrentReads) { + constexpr size_t kBlockSize = 4096; + constexpr size_t kThreadCount = 4; + constexpr size_t kDataSize = kBlockSize * kThreadCount; + + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + std::vector expected(kDataSize); + for (size_t i = 0; i < kThreadCount; ++i) { + std::fill(expected.begin() + i * kBlockSize, + expected.begin() + (i + 1) * kBlockSize, + static_cast(i + 1)); + } + ASSERT_EQ(::pwrite(file.fd(), expected.data(), expected.size(), 0), + static_cast(expected.size())); + + std::array outputs{}; + for (void *&output : outputs) { + DiskAnnUtil::alloc_aligned(&output, kBlockSize, kBlockSize); + ASSERT_NE(output, nullptr); + } + + LinuxAlignedFileReader reader; + reader.open(file.path()); + std::array statuses{}; + std::vector threads; + threads.reserve(kThreadCount); + for (size_t i = 0; i < kThreadCount; ++i) { + threads.emplace_back([&, i]() { + IOContext context{}; + std::vector requests; + requests.emplace_back(i * kBlockSize, kBlockSize, outputs[i]); + statuses[i] = reader.read(requests, context); + }); + } + for (auto &thread : threads) { + thread.join(); + } + + for (size_t i = 0; i < kThreadCount; ++i) { + EXPECT_EQ(statuses[i], 0); + EXPECT_EQ( + std::memcmp(outputs[i], expected.data() + i * kBlockSize, kBlockSize), + 0); + DiskAnnUtil::free_aligned(outputs[i]); + } + reader.close(); +} + +TEST(DiskAnnMobileCompatTest, BuildDumpLoadAndSearch) { + constexpr size_t kDimension = 10; + constexpr size_t kDocCount = 64; + constexpr uint64_t kExpectedKey = 12; + + TemporaryFile index_file; + ASSERT_GE(index_file.fd(), 0); + index_file.release_descriptor_and_unlink(); + + IndexMeta meta(IndexMeta::DataType::DT_FP32, kDimension); + meta.set_metric("SquaredEuclidean", 0, ailego::Params()); + + auto holder = + std::make_shared>( + kDimension); + for (size_t i = 0; i < kDocCount; ++i) { + ailego::NumericalVector vector(kDimension, static_cast(i)); + ASSERT_TRUE(holder->emplace(i, vector)); + } + + ailego::Params build_params; + build_params.set("zvec.diskann.builder.max_degree", 16); + build_params.set("zvec.diskann.builder.list_size", 32); + build_params.set("zvec.diskann.builder.max_pq_chunk_num", 2); + build_params.set("zvec.diskann.builder.threads", 2); + + IndexBuilder::Pointer builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); + ASSERT_NE(builder, nullptr); + ASSERT_EQ(builder->init(meta, build_params), 0); + ASSERT_EQ(builder->train(holder), 0); + ASSERT_EQ(builder->build(holder), 0); + + auto dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_NE(dumper, nullptr); + ASSERT_EQ(dumper->create(index_file.path()), 0); + ASSERT_EQ(builder->dump(dumper), 0); + ASSERT_EQ(dumper->close(), 0); + + int snapshot_fd = ::open(index_file.path(), O_RDONLY); + ASSERT_GE(snapshot_fd, 0); + struct stat snapshot_stat {}; + ASSERT_EQ(::fstat(snapshot_fd, &snapshot_stat), 0); + ASSERT_GT(snapshot_stat.st_size, 4096); + std::vector snapshot(static_cast(snapshot_stat.st_size)); + ASSERT_EQ(::pread(snapshot_fd, snapshot.data(), snapshot.size(), 0), + static_cast(snapshot.size())); + ASSERT_EQ(::close(snapshot_fd), 0); + + IndexSearcher::Pointer searcher = + IndexFactory::CreateSearcher("DiskAnnSearcher"); + ASSERT_NE(searcher, nullptr); + + ailego::Params search_params; + search_params.set("zvec.diskann.searcher.list_size", 64); + ASSERT_EQ(searcher->init(search_params), 0); + + auto storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(storage->open(index_file.path(), false), 0); + ASSERT_EQ(searcher->load(storage, IndexMetric::Pointer()), 0); + + auto context = searcher->create_context(); + ASSERT_NE(context, nullptr); + context->set_topk(5); + + ailego::NumericalVector query(kDimension, 12.1f); + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, kDimension); + ASSERT_EQ(searcher->search_impl(query.data(), query_meta, context), 0); + + const auto &result = context->result(); + ASSERT_FALSE(result.empty()); + EXPECT_NE( + std::find_if(result.begin(), result.end(), + [](const auto &item) { return item.key() == kExpectedKey; }), + result.end()); + + IndexStreamer::Pointer first_streamer = + IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(first_streamer, nullptr); + ASSERT_EQ(first_streamer->init(meta, search_params), 0); + auto first_streamer_storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(first_streamer_storage, nullptr); + ASSERT_EQ(first_streamer_storage->open(index_file.path(), false), 0); + ASSERT_EQ(first_streamer->open(first_streamer_storage), 0); + + IndexStreamer::Pointer second_streamer = + IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(second_streamer, nullptr); + ASSERT_EQ(second_streamer->init(meta, search_params), 0); + auto second_streamer_storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(second_streamer_storage, nullptr); + ASSERT_EQ(second_streamer_storage->open(index_file.path(), false), 0); + ASSERT_EQ(second_streamer->open(second_streamer_storage), 0); + + auto switching_context = first_streamer->create_context(); + ASSERT_NE(switching_context, nullptr); + switching_context->set_topk(5); + switching_context->set_filter( + [](uint64_t key) { return key != kExpectedKey; }); + ASSERT_EQ( + second_streamer->search_impl(query.data(), query_meta, switching_context), + 0); + ASSERT_EQ(switching_context->result().size(), 1u); + EXPECT_EQ(switching_context->result().front().key(), kExpectedKey); + + context.reset(); + searcher.reset(); + storage.reset(); + + ASSERT_EQ(::truncate(index_file.path(), snapshot.size() - 4096), 0); + searcher = IndexFactory::CreateSearcher("DiskAnnSearcher"); + ASSERT_NE(searcher, nullptr); + ASSERT_EQ(searcher->init(search_params), 0); + storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(storage, nullptr); + int corrupt_open_result = storage->open(index_file.path(), false); + bool corrupt_index_rejected = corrupt_open_result != 0; + if (corrupt_open_result == 0) { + corrupt_index_rejected = + searcher->load(storage, IndexMetric::Pointer()) != 0; + } + EXPECT_TRUE(corrupt_index_rejected); + + searcher.reset(); + storage.reset(); + int restore_fd = ::open(index_file.path(), O_WRONLY | O_TRUNC); + ASSERT_GE(restore_fd, 0); + ASSERT_EQ(::pwrite(restore_fd, snapshot.data(), snapshot.size(), 0), + static_cast(snapshot.size())); + ASSERT_EQ(::fsync(restore_fd), 0); + ASSERT_EQ(::close(restore_fd), 0); + + searcher = IndexFactory::CreateSearcher("DiskAnnSearcher"); + ASSERT_NE(searcher, nullptr); + ASSERT_EQ(searcher->init(search_params), 0); + storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(storage->open(index_file.path(), false), 0); + ASSERT_EQ(searcher->load(storage, IndexMetric::Pointer()), 0); + context = searcher->create_context(); + ASSERT_NE(context, nullptr); + context->set_topk(5); + ASSERT_EQ(searcher->search_impl(query.data(), query_meta, context), 0); + EXPECT_NE( + std::find_if(context->result().begin(), context->result().end(), + [](const auto &item) { return item.key() == kExpectedKey; }), + context->result().end()); +} + +} // namespace +} // namespace zvec::core diff --git a/tests/db/CMakeLists.txt b/tests/db/CMakeLists.txt index bfad2bc57..48fa9889e 100644 --- a/tests/db/CMakeLists.txt +++ b/tests/db/CMakeLists.txt @@ -22,10 +22,21 @@ if(APPLE) endif() file(GLOB ALL_TEST_SRCS *_test.cc) + +# The collection DiskAnn stress cases repeatedly rebuild large indexes and are +# intended for desktop CI. Mobile CI exercises DiskAnn through the focused +# diskann_mobile_collection_test target and the core compatibility suite. +if(ANDROID OR IOS) + set(DISKANN_STRESS_TESTS 0) +else() + set(DISKANN_STRESS_TESTS 1) +endif() + foreach(CC_SRCS ${ALL_TEST_SRCS}) get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) cc_gmock( NAME ${CC_TARGET} STRICT + DEFS DISKANN_STRESS_TESTS=${DISKANN_STRESS_TESTS} LIBS zvec zvec_proto core_knn_flat diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index f98ef44f7..e9c5fe78f 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -3088,7 +3088,7 @@ TEST_F(CollectionTest, Feature_Optimize_Repeated) { run_repeated_optimize_test( enable_mmap, std::make_shared( MetricType::IP, 10, 4, false, QuantizeType::FP16)); -#if DISKANN_SUPPORTED +#if DISKANN_SUPPORTED && DISKANN_STRESS_TESTS run_repeated_optimize_test( enable_mmap, std::make_shared( MetricType::IP, 10, 4, 0, QuantizeType::UNDEFINED)); @@ -5628,7 +5628,7 @@ TEST_F(CollectionTest, Feature_Optimize_HNSW_RABITQ) { } #endif -#if DISKANN_SUPPORTED +#if DISKANN_SUPPORTED && DISKANN_STRESS_TESTS TEST_F(CollectionTest, Feature_Optimize_DiskAnn) { auto func = [](MetricType metric_type, int concurrency) { FileHelper::RemoveDirectory(col_path); diff --git a/tests/db/diskann_mobile_collection_test.cc b/tests/db/diskann_mobile_collection_test.cc new file mode 100644 index 000000000..25228a548 --- /dev/null +++ b/tests/db/diskann_mobile_collection_test.cc @@ -0,0 +1,688 @@ +// Copyright 2025-present the zvec project +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(__APPLE__) +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace zvec { +namespace { + +#if defined(__ANDROID__) || \ + (defined(__APPLE__) && (TARGET_OS_IOS || TARGET_OS_SIMULATOR)) +static_assert(DISKANN_SUPPORTED == 1, + "Android and iOS must compile the DiskAnn mobile contract"); +#endif + +#if DISKANN_SUPPORTED + +constexpr char kCollectionPath[] = "diskann_mobile_collection"; +constexpr char kFp32Field[] = "dense_fp32"; +constexpr char kFp16Field[] = "dense_fp16"; +constexpr char kDynamicField[] = "dense_dynamic"; +constexpr char kGroupByField[] = "dense_group_by"; +constexpr size_t kDimension = 16; +constexpr uint64_t kDocCount = 48; + +std::vector MakeFp32Vector(uint64_t doc_id) { + std::vector result(kDimension); + for (size_t i = 0; i < result.size(); ++i) { + result[i] = static_cast(((doc_id + 3) * (i + 5)) % 23) / 23.0F + + static_cast(doc_id) * 0.01F; + } + return result; +} + +std::vector MakeFp16Vector(uint64_t doc_id) { + auto fp32 = MakeFp32Vector(doc_id); + std::vector result; + result.reserve(fp32.size()); + for (float value : fp32) { + result.emplace_back(value); + } + return result; +} + +CollectionSchema::Ptr MakeSchema(MetricType metric, bool include_fp16 = false, + bool include_dynamic = false, + bool include_group_by = false) { + auto schema = std::make_shared("diskann_mobile"); + schema->set_max_doc_count_per_segment(1000); + EXPECT_TRUE(schema + ->add_field(std::make_shared( + "category", DataType::INT32, false)) + .ok()); + EXPECT_TRUE(schema + ->add_field(std::make_shared( + "name", DataType::STRING, false)) + .ok()); + + auto diskann = std::make_shared(metric, 16, 32, 2); + EXPECT_TRUE( + schema + ->add_field(std::make_shared( + kFp32Field, DataType::VECTOR_FP32, kDimension, false, diskann)) + .ok()); + if (include_group_by) { + EXPECT_TRUE(schema + ->add_field(std::make_shared( + kGroupByField, DataType::VECTOR_FP32, kDimension, false, + std::make_shared(metric))) + .ok()); + } + if (include_fp16) { + EXPECT_TRUE(schema + ->add_field(std::make_shared( + kFp16Field, DataType::VECTOR_FP16, kDimension, false, + diskann->clone())) + .ok()); + } + if (include_dynamic) { + EXPECT_TRUE( + schema + ->add_field(std::make_shared( + kDynamicField, DataType::VECTOR_FP32, kDimension, false)) + .ok()); + } + return schema; +} + +Doc MakeDoc(uint64_t doc_id, bool include_fp16 = false, + bool include_dynamic = false, bool include_group_by = false, + std::string pk = "") { + Doc doc; + doc.set_pk(pk.empty() ? "pk_" + std::to_string(doc_id) : std::move(pk)); + doc.set("category", static_cast(doc_id % 4)); + doc.set("name", "name_" + std::to_string(doc_id)); + doc.set>(kFp32Field, MakeFp32Vector(doc_id)); + if (include_group_by) { + doc.set>(kGroupByField, MakeFp32Vector(doc_id)); + } + if (include_fp16) { + doc.set>(kFp16Field, MakeFp16Vector(doc_id)); + } + if (include_dynamic) { + doc.set>(kDynamicField, MakeFp32Vector(doc_id + 7)); + } + return doc; +} + +std::vector MakeDocs(uint64_t begin, uint64_t end, + bool include_fp16 = false, + bool include_dynamic = false, + bool include_group_by = false) { + std::vector docs; + docs.reserve(end - begin); + for (uint64_t doc_id = begin; doc_id < end; ++doc_id) { + docs.emplace_back( + MakeDoc(doc_id, include_fp16, include_dynamic, include_group_by)); + } + return docs; +} + +SearchQuery MakeFp32Query(uint64_t doc_id, const std::string &field, + int topk = 5) { + auto vector = field == kDynamicField ? MakeFp32Vector(doc_id + 7) + : MakeFp32Vector(doc_id); + SearchQuery query; + query.topk_ = topk; + query.target_.field_name_ = field; + query.target_.query_params_ = std::make_shared(32); + query.target_.set_vector( + std::string(reinterpret_cast(vector.data()), + vector.size() * sizeof(float))); + return query; +} + +SearchQuery MakeFp16Query(uint64_t doc_id, int topk = 5) { + auto vector = MakeFp16Vector(doc_id); + SearchQuery query; + query.topk_ = topk; + query.target_.field_name_ = kFp16Field; + query.target_.query_params_ = std::make_shared(32); + query.target_.set_vector( + std::string(reinterpret_cast(vector.data()), + vector.size() * sizeof(float16_t))); + return query; +} + +SearchQuery MakeFlatQuery(uint64_t doc_id, int topk = 5) { + auto vector = MakeFp32Vector(doc_id); + SearchQuery query; + query.topk_ = topk; + query.target_.field_name_ = kGroupByField; + query.target_.query_params_ = std::make_shared(); + query.target_.set_vector( + std::string(reinterpret_cast(vector.data()), + vector.size() * sizeof(float))); + return query; +} + +std::vector SortedPks(const DocPtrList &docs) { + std::vector pks; + pks.reserve(docs.size()); + for (const auto &doc : docs) { + if (doc != nullptr) { + pks.emplace_back(doc->pk()); + } + } + std::sort(pks.begin(), pks.end()); + return pks; +} + +bool FetchContainsPk(const Result &result, const std::string &pk) { + if (!result.has_value() || result->size() != 1) { + return false; + } + auto it = result->find(pk); + return it != result->end() && it->second != nullptr; +} + +::testing::AssertionResult WriteSucceeded(const Result &result, + size_t expected_count) { + if (!result.has_value()) { + return ::testing::AssertionFailure() << result.error().message(); + } + if (result->size() != expected_count) { + return ::testing::AssertionFailure() + << "expected " << expected_count << " write results, got " + << result->size(); + } + for (size_t i = 0; i < result->size(); ++i) { + if (!result->at(i).ok()) { + return ::testing::AssertionFailure() + << "write " << i << " failed: " << result->at(i).message(); + } + } + return ::testing::AssertionSuccess(); +} + +class DiskAnnMobileCollectionTest : public ::testing::Test { + protected: + void SetUp() override { + ailego::FileHelper::RemoveDirectory(kCollectionPath); + } + + void TearDown() override { + ailego::FileHelper::RemoveDirectory(kCollectionPath); + } + + static CollectionOptions Options(bool read_only = false) { + return CollectionOptions{read_only, true, 32 * 1024 * 1024}; + } +}; + +TEST_F(DiskAnnMobileCollectionTest, PublicCollectionApiLifecycle) { + auto schema = MakeSchema(MetricType::L2, false, true); + auto options = Options(); + auto create_result = + Collection::CreateAndOpen(kCollectionPath, *schema, options); + ASSERT_TRUE(create_result.has_value()) << create_result.error().message(); + auto collection = std::move(create_result.value()); + + auto path_result = collection->Path(); + ASSERT_TRUE(path_result.has_value()) << path_result.error().message(); + EXPECT_EQ(*path_result, kCollectionPath); + auto schema_result = collection->Schema(); + ASSERT_TRUE(schema_result.has_value()) << schema_result.error().message(); + EXPECT_EQ(*schema_result, *schema); + auto options_result = collection->Options(); + ASSERT_TRUE(options_result.has_value()) << options_result.error().message(); + EXPECT_EQ(*options_result, options); + auto empty_stats = collection->Stats(); + ASSERT_TRUE(empty_stats.has_value()) << empty_stats.error().message(); + EXPECT_EQ(empty_stats->doc_count, 0u); + + auto docs = MakeDocs(0, 32, false, true); + ASSERT_TRUE(WriteSucceeded(collection->Insert(docs), docs.size())); + ASSERT_TRUE(collection->Flush().ok()); + auto flushed_stats = collection->Stats(); + ASSERT_TRUE(flushed_stats.has_value()) << flushed_stats.error().message(); + ASSERT_EQ(flushed_stats->index_completeness[kFp32Field], 0); + ASSERT_TRUE(collection->Optimize(OptimizeOptions{2}).ok()); + auto optimized_stats = collection->Stats(); + ASSERT_TRUE(optimized_stats.has_value()) << optimized_stats.error().message(); + ASSERT_EQ(optimized_stats->index_completeness[kFp32Field], 1); + + auto fetch = collection->Fetch( + {"pk_8"}, std::vector{"category", "name"}, false); + ASSERT_TRUE(fetch.has_value()) << fetch.error().message(); + ASSERT_TRUE(FetchContainsPk(fetch, "pk_8")); + EXPECT_TRUE(fetch->at("pk_8")->has("category")); + EXPECT_TRUE(fetch->at("pk_8")->has("name")); + EXPECT_FALSE(fetch->at("pk_8")->has(kFp32Field)); + + std::vector update_docs{MakeDoc(100, false, true, false, "pk_0")}; + ASSERT_TRUE( + WriteSucceeded(collection->Update(update_docs), update_docs.size())); + + std::vector upsert_docs{MakeDoc(101, false, true, false, "pk_1"), + MakeDoc(32, false, true)}; + ASSERT_TRUE( + WriteSucceeded(collection->Upsert(upsert_docs), upsert_docs.size())); + ASSERT_TRUE(WriteSucceeded(collection->Delete({"pk_2"}), 1)); + ASSERT_TRUE(collection->DeleteByFilter("category = 3").ok()); + auto deleted_fetch = collection->Fetch({"pk_2", "pk_3"}); + ASSERT_TRUE(deleted_fetch.has_value()) << deleted_fetch.error().message(); + ASSERT_EQ(deleted_fetch->size(), 2u); + auto deleted_pk2 = deleted_fetch->find("pk_2"); + auto deleted_pk3 = deleted_fetch->find("pk_3"); + ASSERT_NE(deleted_pk2, deleted_fetch->end()); + ASSERT_NE(deleted_pk3, deleted_fetch->end()); + EXPECT_EQ(deleted_pk2->second, nullptr); + EXPECT_EQ(deleted_pk3->second, nullptr); + + auto added_field = + std::make_shared("category_copy", DataType::INT32, false); + ASSERT_TRUE(collection->AddColumn(added_field, "category").ok()); + ASSERT_TRUE( + collection->AlterColumn("category_copy", "category_renamed").ok()); + ASSERT_TRUE(collection->DropColumn("category_renamed").ok()); + + auto dynamic_index = + std::make_shared(MetricType::L2, 16, 32, 2); + ASSERT_TRUE(collection->CreateIndex(kDynamicField, dynamic_index).ok()); + ASSERT_TRUE(collection->Optimize(OptimizeOptions{2}).ok()); + auto dynamic_result = collection->Query(MakeFp32Query(8, kDynamicField, 32)); + ASSERT_TRUE(dynamic_result.has_value()) << dynamic_result.error().message(); + ASSERT_FALSE(dynamic_result->empty()); + ASSERT_TRUE(collection->DropIndex(kDynamicField).ok()); + + ASSERT_TRUE(collection->Flush().ok()); + collection.reset(); + + auto reopen_result = Collection::Open(kCollectionPath, options); + ASSERT_TRUE(reopen_result.has_value()) << reopen_result.error().message(); + collection = std::move(reopen_result.value()); + auto primary_result = collection->Query(MakeFp32Query(8, kFp32Field)); + ASSERT_TRUE(primary_result.has_value()) << primary_result.error().message(); + ASSERT_FALSE(primary_result->empty()); + auto reopened_stats = collection->Stats(); + ASSERT_TRUE(reopened_stats.has_value()) << reopened_stats.error().message(); + EXPECT_LT(reopened_stats->doc_count, 33u); + collection.reset(); + + auto read_only_result = Collection::Open(kCollectionPath, Options(true)); + ASSERT_TRUE(read_only_result.has_value()) + << read_only_result.error().message(); + collection = std::move(read_only_result.value()); + auto read_only_query = collection->Query(MakeFp32Query(8, kFp32Field)); + ASSERT_TRUE(read_only_query.has_value()) << read_only_query.error().message(); + ASSERT_FALSE(read_only_query->empty()); + auto rejected_docs = MakeDocs(40, 41, false, true); + EXPECT_FALSE(collection->Insert(rejected_docs).has_value()); + EXPECT_FALSE(collection->Optimize().ok()); + collection.reset(); + + reopen_result = Collection::Open(kCollectionPath, options); + ASSERT_TRUE(reopen_result.has_value()) << reopen_result.error().message(); + collection = std::move(reopen_result.value()); + ASSERT_TRUE(collection->Destroy().ok()); + EXPECT_FALSE(collection->Stats().has_value()); + EXPECT_FALSE(Collection::Open(kCollectionPath, options).has_value()); +} + +TEST_F(DiskAnnMobileCollectionTest, CompleteQuerySurfaceAndMetricMatrix) { + for (MetricType metric : + {MetricType::L2, MetricType::IP, MetricType::COSINE}) { + SCOPED_TRACE(static_cast(metric)); + ailego::FileHelper::RemoveDirectory(kCollectionPath); + auto schema = MakeSchema(metric, true, false, true); + auto create_result = + Collection::CreateAndOpen(kCollectionPath, *schema, Options()); + ASSERT_TRUE(create_result.has_value()) << create_result.error().message(); + auto collection = std::move(create_result.value()); + + auto docs = MakeDocs(0, kDocCount, true, false, true); + ASSERT_TRUE(WriteSucceeded(collection->Insert(docs), docs.size())); + ASSERT_TRUE(collection->Flush().ok()); + ASSERT_TRUE(collection->Optimize(OptimizeOptions{2}).ok()); + + auto fp32_query = MakeFp32Query(12, kFp32Field, 8); + fp32_query.filter_ = "category = 0"; + fp32_query.include_vector_ = true; + fp32_query.include_doc_id_ = true; + fp32_query.output_fields_ = std::vector{"category", "name"}; + auto fp32_result = collection->Query(fp32_query); + ASSERT_TRUE(fp32_result.has_value()) << fp32_result.error().message(); + ASSERT_FALSE(fp32_result->empty()); + for (const auto &doc : *fp32_result) { + ASSERT_NE(doc, nullptr); + auto category = doc->get("category"); + ASSERT_TRUE(category.has_value()); + EXPECT_EQ(category.value(), 0); + EXPECT_TRUE(doc->has("name")); + EXPECT_TRUE(doc->has(kFp32Field)); + } + EXPECT_TRUE(std::any_of(fp32_result->begin(), fp32_result->end(), + [](const Doc::Ptr &doc) { + return doc != nullptr && doc->doc_id() != 0; + })); + auto default_params_query = MakeFp32Query(12, kFp32Field); + default_params_query.target_.query_params_.reset(); + auto default_params_result = collection->Query(default_params_query); + ASSERT_TRUE(default_params_result.has_value()) + << default_params_result.error().message(); + ASSERT_FALSE(default_params_result->empty()); + + SearchQuery scalar_query; + scalar_query.topk_ = 5; + scalar_query.filter_ = "category = 1"; + scalar_query.output_fields_ = std::vector{"category", "name"}; + auto scalar_result = collection->Query(scalar_query); + ASSERT_TRUE(scalar_result.has_value()) << scalar_result.error().message(); + ASSERT_EQ(scalar_result->size(), 5u); + for (const auto &doc : *scalar_result) { + ASSERT_NE(doc, nullptr); + auto category = doc->get("category"); + ASSERT_TRUE(category.has_value()); + EXPECT_EQ(category.value(), 1); + } + + auto fp16_result = collection->Query(MakeFp16Query(12, 8)); + ASSERT_TRUE(fp16_result.has_value()) << fp16_result.error().message(); + ASSERT_FALSE(fp16_result->empty()); + + ASSERT_GE(fp32_result->size(), 2u); + const float best_score = fp32_result->front()->score(); + const float worst_score = fp32_result->back()->score(); + const float radius = (best_score + worst_score) / 2.0F; + ASSERT_GT(radius, 0.0F); + auto radius_query = MakeFp32Query(12, kFp32Field, 8); + radius_query.filter_ = "category = 0"; + radius_query.target_.query_params_->set_radius(radius); + auto radius_result = collection->Query(radius_query); + ASSERT_TRUE(radius_result.has_value()) << radius_result.error().message(); + ASSERT_FALSE(radius_result->empty()); + EXPECT_LT(radius_result->size(), fp32_result->size()); + for (const auto &doc : *radius_result) { + ASSERT_NE(doc, nullptr); + if (metric == MetricType::IP) { + EXPECT_GE(doc->score(), radius); + } else { + EXPECT_LE(doc->score(), radius); + } + } + + MultiQuery multi_query; + multi_query.topk = 8; + multi_query.filter = "category = 0"; + multi_query.include_vector = true; + multi_query.include_doc_id_ = true; + multi_query.output_fields = std::vector{"category", "name"}; + multi_query.rerank = reranker::RrfParams{60}; + for (uint64_t doc_id : {12u, 20u}) { + auto search_query = MakeFp32Query(doc_id, kFp32Field, 16); + SubQuery sub_query; + sub_query.target_ = std::move(search_query.target_); + sub_query.num_candidates_ = 16; + multi_query.queries.emplace_back(std::move(sub_query)); + } + auto multi_result = collection->Query(multi_query); + ASSERT_TRUE(multi_result.has_value()) << multi_result.error().message(); + ASSERT_FALSE(multi_result->empty()); + EXPECT_LE(multi_result->size(), 8u); + for (const auto &doc : *multi_result) { + ASSERT_NE(doc, nullptr); + EXPECT_TRUE(doc->has("category")); + EXPECT_TRUE(doc->has("name")); + EXPECT_TRUE(doc->has(kFp32Field)); + auto category = doc->get("category"); + ASSERT_TRUE(category.has_value()); + EXPECT_EQ(category.value(), 0); + } + EXPECT_TRUE(std::any_of(multi_result->begin(), multi_result->end(), + [](const Doc::Ptr &doc) { + return doc != nullptr && doc->doc_id() != 0; + })); + + GroupByVectorQuery group_query; + group_query.target_ = MakeFlatQuery(12, 8).target_; + group_query.filter_ = "category >= 0"; + group_query.group_by_field_name_ = "category"; + group_query.group_count_ = 4; + group_query.topk_per_group_ = 2; + group_query.include_vector_ = true; + group_query.output_fields_ = std::vector{"category", "name"}; + auto group_result = collection->GroupByQuery(group_query); + ASSERT_TRUE(group_result.has_value()) << group_result.error().message(); + ASSERT_FALSE(group_result->empty()); + EXPECT_LE(group_result->size(), 4u); + for (const auto &group : *group_result) { + EXPECT_FALSE(group.group_by_value_.empty()); + EXPECT_FALSE(group.docs_.empty()); + EXPECT_LE(group.docs_.size(), 2u); + for (const auto &doc : group.docs_) { + EXPECT_TRUE(doc.has("category")); + EXPECT_TRUE(doc.has("name")); + EXPECT_TRUE(doc.has(kGroupByField)); + } + } + + auto selected_fetch = collection->Fetch( + {"pk_12"}, std::vector{"category"}, false); + ASSERT_TRUE(selected_fetch.has_value()) << selected_fetch.error().message(); + ASSERT_TRUE(FetchContainsPk(selected_fetch, "pk_12")); + EXPECT_TRUE(selected_fetch->at("pk_12")->has("category")); + EXPECT_FALSE(selected_fetch->at("pk_12")->has("name")); + EXPECT_FALSE(selected_fetch->at("pk_12")->has(kFp32Field)); + + collection.reset(); + auto reopen_result = Collection::Open(kCollectionPath, Options()); + ASSERT_TRUE(reopen_result.has_value()) << reopen_result.error().message(); + collection = std::move(reopen_result.value()); + auto reopened_query = collection->Query(MakeFp32Query(12, kFp32Field)); + ASSERT_TRUE(reopened_query.has_value()) << reopened_query.error().message(); + ASSERT_FALSE(reopened_query->empty()); + } +} + +TEST_F(DiskAnnMobileCollectionTest, ConcurrentQueryAndFetch) { + constexpr size_t kThreadCount = 4; + constexpr size_t kIterations = 20; + + auto schema = MakeSchema(MetricType::L2); + auto create_result = + Collection::CreateAndOpen(kCollectionPath, *schema, Options()); + ASSERT_TRUE(create_result.has_value()) << create_result.error().message(); + auto collection = std::move(create_result.value()); + auto docs = MakeDocs(0, kDocCount); + ASSERT_TRUE(WriteSucceeded(collection->Insert(docs), docs.size())); + ASSERT_TRUE(collection->Optimize(OptimizeOptions{2}).ok()); + + std::array, kDocCount> query_baselines; + for (uint64_t doc_id = 0; doc_id < kDocCount; ++doc_id) { + auto query_result = collection->Query(MakeFp32Query(doc_id, kFp32Field)); + ASSERT_TRUE(query_result.has_value()) << query_result.error().message(); + ASSERT_FALSE(query_result->empty()); + query_baselines[doc_id] = SortedPks(*query_result); + ASSERT_EQ(query_baselines[doc_id].size(), query_result->size()); + } + + std::atomic failure_count{0}; + std::mutex failure_mutex; + std::vector failures; + auto record_failure = [&](const std::string &failure) { + ++failure_count; + std::lock_guard lock(failure_mutex); + failures.emplace_back(failure); + }; + std::vector threads; + threads.reserve(kThreadCount); + for (size_t thread_id = 0; thread_id < kThreadCount; ++thread_id) { + threads.emplace_back([&, thread_id]() { + for (size_t iteration = 0; iteration < kIterations; ++iteration) { + uint64_t doc_id = (thread_id * kIterations + iteration) % kDocCount; + auto query_result = + collection->Query(MakeFp32Query(doc_id, kFp32Field)); + auto fetch_result = collection->Fetch({"pk_" + std::to_string(doc_id)}); + bool query_ok = query_result.has_value() && + SortedPks(*query_result) == query_baselines[doc_id]; + bool fetch_ok = + FetchContainsPk(fetch_result, "pk_" + std::to_string(doc_id)); + if (!query_ok || !fetch_ok) { + record_failure( + "shared collection: thread=" + std::to_string(thread_id) + + ", iteration=" + std::to_string(iteration) + + ", query_ok=" + std::to_string(query_ok) + + ", fetch_ok=" + std::to_string(fetch_ok)); + } + } + }); + } + for (auto &thread : threads) { + thread.join(); + } + + EXPECT_EQ(failure_count.load(), 0u); + EXPECT_TRUE(failures.empty()) << (failures.empty() ? "" : failures.front()); + + ASSERT_TRUE(collection->Flush().ok()); + collection.reset(); + failure_count.store(0); + failures.clear(); + threads.clear(); + for (size_t thread_id = 0; thread_id < kThreadCount; ++thread_id) { + threads.emplace_back([&, thread_id]() { + auto open_result = Collection::Open(kCollectionPath, Options(true)); + if (!open_result.has_value()) { + record_failure("read-only open: thread=" + std::to_string(thread_id)); + return; + } + auto read_only_collection = std::move(open_result.value()); + for (size_t iteration = 0; iteration < kIterations; ++iteration) { + uint64_t doc_id = (thread_id * kIterations + iteration) % kDocCount; + auto query_result = + read_only_collection->Query(MakeFp32Query(doc_id, kFp32Field)); + auto fetch_result = + read_only_collection->Fetch({"pk_" + std::to_string(doc_id)}); + bool query_ok = query_result.has_value() && + SortedPks(*query_result) == query_baselines[doc_id]; + bool fetch_ok = + FetchContainsPk(fetch_result, "pk_" + std::to_string(doc_id)); + if (!query_ok || !fetch_ok) { + record_failure( + "read-only collection: thread=" + std::to_string(thread_id) + + ", iteration=" + std::to_string(iteration) + + ", query_ok=" + std::to_string(query_ok) + + ", fetch_ok=" + std::to_string(fetch_ok)); + } + } + }); + } + for (auto &thread : threads) { + thread.join(); + } + + EXPECT_EQ(failure_count.load(), 0u); + EXPECT_TRUE(failures.empty()) << (failures.empty() ? "" : failures.front()); +} + +TEST_F(DiskAnnMobileCollectionTest, OperationFailuresDoNotPoisonCollection) { + auto schema = MakeSchema(MetricType::L2); + auto create_result = + Collection::CreateAndOpen(kCollectionPath, *schema, Options()); + ASSERT_TRUE(create_result.has_value()) << create_result.error().message(); + auto collection = std::move(create_result.value()); + auto docs = MakeDocs(0, kDocCount); + ASSERT_TRUE(WriteSucceeded(collection->Insert(docs), docs.size())); + ASSERT_TRUE(collection->Optimize(OptimizeOptions{2}).ok()); + + auto invalid_query = MakeFp32Query(12, kFp32Field); + invalid_query.target_.set_vector("invalid-size"); + EXPECT_FALSE(collection->Query(invalid_query).has_value()); + + auto wrong_params_query = MakeFp32Query(12, kFp32Field); + wrong_params_query.target_.query_params_ = + std::make_shared(); + EXPECT_FALSE(collection->Query(wrong_params_query).has_value()); + + Doc invalid_doc; + invalid_doc.set_pk("invalid_doc"); + invalid_doc.set("category", 0); + invalid_doc.set("name", "missing required vector"); + std::vector invalid_docs{invalid_doc}; + auto invalid_write = collection->Insert(invalid_docs); + EXPECT_TRUE(!invalid_write.has_value() || invalid_write->empty() || + !invalid_write->front().ok()); + + GroupByVectorQuery unsupported_group_query; + unsupported_group_query.target_ = MakeFp32Query(12, kFp32Field, 8).target_; + unsupported_group_query.group_by_field_name_ = "category"; + unsupported_group_query.group_count_ = 4; + unsupported_group_query.topk_per_group_ = 2; + EXPECT_FALSE(collection->GroupByQuery(unsupported_group_query).has_value()); + + auto valid_result = collection->Query(MakeFp32Query(12, kFp32Field)); + ASSERT_TRUE(valid_result.has_value()) << valid_result.error().message(); + ASSERT_FALSE(valid_result->empty()); + + auto recovery_docs = MakeDocs(kDocCount, kDocCount + 1); + ASSERT_TRUE( + WriteSucceeded(collection->Insert(recovery_docs), recovery_docs.size())); + ASSERT_TRUE(collection->Flush().ok()); + ASSERT_TRUE(collection->Optimize(OptimizeOptions{2}).ok()); + collection.reset(); + + auto reopen_result = Collection::Open(kCollectionPath, Options()); + ASSERT_TRUE(reopen_result.has_value()) << reopen_result.error().message(); + collection = std::move(reopen_result.value()); + auto recovered_result = + collection->Query(MakeFp32Query(kDocCount, kFp32Field)); + ASSERT_TRUE(recovered_result.has_value()) + << recovered_result.error().message(); + ASSERT_FALSE(recovered_result->empty()); + const std::string recovered_pk = "pk_" + std::to_string(kDocCount); + auto recovered_fetch = collection->Fetch({recovered_pk}); + EXPECT_TRUE(FetchContainsPk(recovered_fetch, recovered_pk)); + auto recovered_stats = collection->Stats(); + ASSERT_TRUE(recovered_stats.has_value()) << recovered_stats.error().message(); + EXPECT_EQ(recovered_stats->doc_count, kDocCount + 1); +} + +#else + +TEST(DiskAnnMobileCollectionTest, PlatformDoesNotClaimMobileSupport) { + GTEST_SKIP() << "DiskAnn is not enabled on this desktop platform"; +} + +#endif + +} // namespace +} // namespace zvec