diff --git a/.github/cmake/sdk-smoke/CMakeLists.txt b/.github/cmake/sdk-smoke/CMakeLists.txt index a46d59a7d..9ef975cf4 100644 --- a/.github/cmake/sdk-smoke/CMakeLists.txt +++ b/.github/cmake/sdk-smoke/CMakeLists.txt @@ -18,6 +18,10 @@ file(GLOB _smoke_cpp_headers "${ZVEC_SDK_DIR}/include/zvec/db/*.h") if(NOT _smoke_cpp_headers) message(FATAL_ERROR "SDK package is missing C++ headers under include/zvec/db/") endif() +if(ANDROID AND NOT EXISTS "${ZVEC_SDK_DIR}/lib/libc++_shared.so") + message(FATAL_ERROR + "Android C++ SDK is missing its required lib/libc++_shared.so runtime") +endif() # The jieba FTS dict must ship with the SDK (mirrors the Python wheel's # zvec/data/jieba_dict bundle). Consumers register the directory via # zvec_set_default_jieba_dict_dir() or ZVEC_JIEBA_DICT_DIR. diff --git a/.github/cmake/sdk-smoke/smoke.cc b/.github/cmake/sdk-smoke/smoke.cc index cd8dbc0c2..62493cf9d 100644 --- a/.github/cmake/sdk-smoke/smoke.cc +++ b/.github/cmake/sdk-smoke/smoke.cc @@ -1,4 +1,5 @@ #include +#include #include int main() { @@ -7,7 +8,17 @@ int main() { std::cerr << "cpp_smoke: GetDefaultMessage() returned null" << std::endl; return 1; } + const auto io_backend_type = zvec::ailego::current_io_backend_type(); + const std::string io_backend_description = + zvec::ailego::current_io_backend_description(); + if (io_backend_description.empty()) { + std::cerr << "cpp_smoke: current_io_backend_description() returned empty" + << std::endl; + return 1; + } std::cout << "cpp_smoke: StatusCode::OK -> " << message << std::endl; + std::cout << "cpp_smoke: I/O backend " << static_cast(io_backend_type) + << " -> " << io_backend_description << std::endl; std::cout << "cpp_smoke: OK" << std::endl; return 0; } diff --git a/.github/workflows/04-android-build.yml b/.github/workflows/04-android-build.yml index 6b2abd5fb..e3bd914ca 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 || \ @@ -92,8 +100,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 \ @@ -101,8 +113,19 @@ 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 + cmake --build "$BUILD_DIR" \ + --target diskann_mobile_compat_test diskann_mobile_collection_test \ + --parallel # Discover test targets from ctest metadata echo "Discovering test targets..." @@ -121,11 +144,56 @@ jobs: done < <(ninja -C "$BUILD_DIR" -t targets all 2>/dev/null || true) fi + if [ ${#TEST_NAMES[@]} -eq 0 ]; then + echo "ERROR: No Android test targets were discovered" + exit 1 + fi + + # All test cases in this legacy target are disabled with #if 0. + # Exclude it instead of treating a zero-test binary as a passing suite. + FILTERED_TEST_NAMES=() + for name in "${TEST_NAMES[@]}"; do + if [ "$name" = "cosine_distance_matrix_int8_test" ]; then + echo "Skipping $name: it contains no enabled GTest cases" + continue + fi + FILTERED_TEST_NAMES+=("$name") + done + TEST_NAMES=("${FILTERED_TEST_NAMES[@]}") + + if [ ${#TEST_NAMES[@]} -eq 0 ]; then + echo "ERROR: No non-empty Android test targets were discovered" + exit 1 + fi + + for required_test in diskann_mobile_compat_test diskann_mobile_collection_test; do + found=0 + for name in "${TEST_NAMES[@]}"; do + if [ "$name" = "$required_test" ]; then + found=1 + break + fi + done + if [ "$found" -ne 1 ]; then + echo "ERROR: Required Android DiskAnn test target was not discovered: $required_test" + exit 1 + fi + done + echo "Building ${#TEST_NAMES[@]} test executables..." ninja -C "$BUILD_DIR" -j$(nproc) "${TEST_NAMES[@]}" - # ── Step 2: start emulator ───────────────────────────────────────── - - name: 'Step 2: Start Android emulator' + for required_test in diskann_mobile_compat_test diskann_mobile_collection_test; do + required_binary="$BUILD_DIR/bin/$required_test" + if [ ! -x "$required_binary" ]; then + echo "ERROR: Missing required Android DiskAnn test binary: $required_binary" + exit 1 + fi + done + + # ── Step 3: start emulator ───────────────────────────────────────── + - name: 'Step 3: Start Android emulator' + if: matrix.abi == 'x86_64' shell: bash run: | AVD_NAME="zvec_test_avd" @@ -218,8 +286,9 @@ jobs: echo "Device ABI: $(adb shell getprop ro.product.cpu.abi | tr -d '\r')" echo "ABI list : $(adb shell getprop ro.product.cpu.abilist | tr -d '\r')" - # ── Step 3: run unit tests on emulator ───────────────────────────── - - name: 'Step 3: Run unit tests on emulator' + # ── 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 }} @@ -262,6 +331,34 @@ jobs: done < <(ninja -C "$BUILD_DIR" -t targets all 2>/dev/null || true) fi + if [ ${#TEST_NAMES[@]} -eq 0 ]; then + echo "ERROR: No Android test targets were discovered" + exit 1 + fi + + FILTERED_TEST_NAMES=() + for name in "${TEST_NAMES[@]}"; do + if [ "$name" = "cosine_distance_matrix_int8_test" ]; then + continue + fi + FILTERED_TEST_NAMES+=("$name") + done + TEST_NAMES=("${FILTERED_TEST_NAMES[@]}") + + for required_test in diskann_mobile_compat_test diskann_mobile_collection_test; do + found=0 + for name in "${TEST_NAMES[@]}"; do + if [ "$name" = "$required_test" ]; then + found=1 + break + fi + done + if [ "$found" -ne 1 ]; then + echo "ERROR: Required Android DiskAnn test target was not discovered: $required_test" + exit 1 + fi + done + # Collect test binaries TEST_BINS=() for name in "${TEST_NAMES[@]}"; do @@ -269,15 +366,31 @@ jobs: if [ -n "$bin_path" ]; then TEST_BINS+=("$bin_path") else - echo "WARNING: binary not found for '$name'" + echo "ERROR: binary not found for discovered test target '$name'" + exit 1 + fi + done + + for required_test in diskann_mobile_compat_test diskann_mobile_collection_test; do + required_binary="$BUILD_DIR/bin/$required_test" + if [ ! -x "$required_binary" ]; then + echo "ERROR: Missing required Android DiskAnn test binary: $required_binary" + exit 1 fi done TOTAL=${#TEST_BINS[@]} + if [ "$TOTAL" -eq 0 ]; then + echo "ERROR: No Android test binaries were collected" + exit 1 + fi + PASSED=0 FAILED=0 FAILED_NAMES=() IDX=0 + DISKANN_COMPAT_PASSED=0 + DISKANN_COLLECTION_PASSED=0 echo "Running $TOTAL unit tests on emulator..." @@ -293,45 +406,90 @@ jobs: echo " [$IDX/$TOTAL] $test_name" echo "────────────────────────────────────────" - set +e # Create isolated working directory adb shell "mkdir -p $WORK_DIR" 2>/dev/null # Copy helper binaries into working directory so crash_recovery tests # (which fork+exec data_generator / collection_optimizer) can find them - adb shell "for h in $DEVICE_TEST_DIR/data_generator $DEVICE_TEST_DIR/collection_optimizer; do [ -f \$h ] && cp \$h $WORK_DIR/; done" 2>/dev/null + adb shell "for h in $DEVICE_TEST_DIR/data_generator $DEVICE_TEST_DIR/collection_optimizer; do [ -f \$h ] && cp \$h $WORK_DIR/; done" 2>/dev/null || true # Push test binary adb push "$test_bin" "$device_path" > /dev/null 2>&1 adb shell "chmod 755 $device_path" 2>/dev/null - # Run test from its own working directory with LD_LIBRARY_PATH - OUTPUT=$(adb shell "cd $WORK_DIR && LD_LIBRARY_PATH=$DEVICE_LIB_DIR $device_path 2>&1; echo EXIT_CODE=\$?" 2>&1) + # Run test from its own working directory with a hard timeout. + # GTEST_COLOR=no keeps the required summary machine-readable. + set +e + OUTPUT=$(timeout --signal=TERM --kill-after=10s 600s \ + adb shell "cd $WORK_DIR && GTEST_COLOR=no LD_LIBRARY_PATH=$DEVICE_LIB_DIR $device_path 2>&1; echo EXIT_CODE=\$?" 2>&1) + ADB_EXIT=$? + set -e # Extract exit code from the output - EXIT_CODE=$(echo "$OUTPUT" | grep -o 'EXIT_CODE=[0-9]*' | tail -1 | cut -d= -f2) - set -e + EXIT_CODE=$(printf '%s\n' "$OUTPUT" | grep -o 'EXIT_CODE=[0-9]*' | tail -1 | cut -d= -f2 || true) # Print test output (without the EXIT_CODE marker) echo "$OUTPUT" | grep -v 'EXIT_CODE=' | sed 's/^/ /' || true - if [ "$EXIT_CODE" = "0" ]; then + EXPECTED_COUNT="" + case "$test_name" in + diskann_mobile_compat_test) + EXPECTED_COUNT=10 + ;; + diskann_mobile_collection_test) + EXPECTED_COUNT=4 + ;; + esac + + TEST_PASSED=1 + if [ "$ADB_EXIT" -eq 124 ] || [ "$ADB_EXIT" -eq 137 ]; then + echo " >>> FAILED (timed out after 600 seconds)" + adb shell "pkill -f $device_path" 2>/dev/null || true + TEST_PASSED=0 + elif [ "$ADB_EXIT" -ne 0 ]; then + echo " >>> FAILED (adb exit=$ADB_EXIT)" + TEST_PASSED=0 + elif [ -z "$EXIT_CODE" ]; then + echo " >>> FAILED (missing device exit code)" + TEST_PASSED=0 + elif [ "$EXIT_CODE" -ne 0 ]; then + echo " >>> FAILED (exit=$EXIT_CODE)" + TEST_PASSED=0 + elif grep -Eq '^\[ FAILED \]' <<< "$OUTPUT"; then + echo " >>> FAILED (GTest reported failures)" + TEST_PASSED=0 + elif [ -n "$EXPECTED_COUNT" ] && \ + ! grep -Eq "^\\[ PASSED \\][[:space:]]+${EXPECTED_COUNT} tests?\\.[[:space:]]*$" <<< "$OUTPUT"; then + echo " >>> FAILED (required ${EXPECTED_COUNT}/${EXPECTED_COUNT} summary not found)" + TEST_PASSED=0 + elif [ -n "$EXPECTED_COUNT" ] && \ + ! grep -Eq "^\\[==========\\][[:space:]]+${EXPECTED_COUNT} tests?[[:space:]]+from[[:space:]]+[1-9][0-9]* test (suites?|cases?)[[:space:]]+ran\\." <<< "$OUTPUT"; then + echo " >>> FAILED (expected exactly ${EXPECTED_COUNT} executed tests)" + TEST_PASSED=0 + elif [ -n "$EXPECTED_COUNT" ]; then + : # The required DiskAnn suite ran exactly N/N tests successfully. + elif grep -Eq '^\[==========\][[:space:]]+[1-9][0-9]* tests?[[:space:]]+from[[:space:]]+[1-9][0-9]* test (suites?|cases?)[[:space:]]+ran\.' <<< "$OUTPUT"; then + : # A non-empty GTest suite completed without failures; skips are valid. + elif [ "$test_name" = "c_api_test" ] && \ + grep -Eq '^Passed: [1-9][0-9]*[[:space:]]*$' <<< "$OUTPUT" && \ + grep -Eq '^Failed: 0[[:space:]]*$' <<< "$OUTPUT"; then + : # c_api_test uses a custom test framework. + else + echo " >>> FAILED (no recognisable non-empty passing test summary)" + TEST_PASSED=0 + fi + + if [ "$TEST_PASSED" -eq 1 ]; then echo " >>> PASSED" PASSED=$((PASSED + 1)) - else - # Detect "crash-on-exit" pattern: all gtest assertions passed but - # process crashed during static destructor teardown (common with c++_static STL) - GTEST_PASSED_LINE=$(echo "$OUTPUT" | grep '\[ PASSED \]' | tail -1 || true) - GTEST_FAILED_LINE=$(echo "$OUTPUT" | grep '\[ FAILED \]' | head -1 || true) - if [ -n "$GTEST_PASSED_LINE" ] && [ -z "$GTEST_FAILED_LINE" ] && \ - { [ "$EXIT_CODE" = "139" ] || [ "$EXIT_CODE" = "134" ] || [ "$EXIT_CODE" = "135" ]; }; then - echo " >>> PASSED (crash-on-exit ignored, exit=$EXIT_CODE)" - PASSED=$((PASSED + 1)) - else - echo " >>> FAILED (exit=$EXIT_CODE)" - FAILED=$((FAILED + 1)) - FAILED_NAMES+=("$test_name") + if [ "$test_name" = "diskann_mobile_compat_test" ]; then + DISKANN_COMPAT_PASSED=1 + elif [ "$test_name" = "diskann_mobile_collection_test" ]; then + DISKANN_COLLECTION_PASSED=1 fi + else + FAILED=$((FAILED + 1)) + FAILED_NAMES+=("$test_name") fi # Clean up binary and working directory to reclaim disk space @@ -354,54 +512,88 @@ jobs: fi echo "============================================================" + if [ "$DISKANN_COMPAT_PASSED" -ne 1 ] || [ "$DISKANN_COLLECTION_PASSED" -ne 1 ]; then + echo "Required DiskAnn tests did not both complete successfully" + exit 1 + fi + if [ "$PASSED" -ne "$TOTAL" ]; then + echo "Only ${PASSED}/${TOTAL} Android tests completed successfully" + exit 1 + fi if [ $FAILED -gt 0 ]; then exit 1 fi echo "All tests passed!" - # ── Step 4: build and run examples ───────────────────────────────── - - name: 'Step 4: Build and run examples' + # ── 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 3; push again in - # case Step 3 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 f5b95974b..854bb660b 100644 --- a/.github/workflows/06-ios-build.yml +++ b/.github/workflows/06-ios-build.yml @@ -10,6 +10,7 @@ permissions: jobs: build-ios: runs-on: macos-15 + timeout-minutes: 120 strategy: fail-fast: false matrix: @@ -50,8 +51,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" \ -DIOS=ON \ @@ -61,6 +66,36 @@ 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 required DiskAnn test targets + run: | + NPROC=$(sysctl -n hw.ncpu) + BUILD_DIR=build_ios_${{ matrix.platform }} + cmake --build "$BUILD_DIR" \ + --target diskann_mobile_compat_test diskann_mobile_collection_test \ + --parallel "$NPROC" + + for test_name in diskann_mobile_compat_test diskann_mobile_collection_test; do + test_app="$BUILD_DIR/bin/${test_name}.app" + if [ ! -d "$test_app" ]; then + echo "::error::Missing required iOS DiskAnn test app: $test_app" + exit 1 + fi + done + - name: Build test targets if: matrix.test_on_simulator run: | @@ -70,6 +105,11 @@ jobs: - name: Boot iOS Simulator if: matrix.test_on_simulator run: | + if [ "$(uname -m)" != "arm64" ]; then + echo "::error::SIMULATORARM64 tests require an arm64 macOS runner" + exit 1 + fi + DEVICE_ID=$(xcrun simctl list devices available -j \ | python3 -c " import json, sys @@ -84,6 +124,7 @@ jobs: ") echo "DEVICE_ID=$DEVICE_ID" >> $GITHUB_ENV xcrun simctl boot "$DEVICE_ID" + xcrun simctl bootstatus "$DEVICE_ID" -b echo "Booted simulator: $DEVICE_ID" - name: Run all tests on simulator @@ -92,52 +133,167 @@ jobs: FAILED_TESTS="" PASSED=0 TOTAL=0 + DISKANN_COMPAT_PASSED=0 + DISKANN_COLLECTION_PASSED=0 + BUILD_DIR=build_ios_${{ matrix.platform }} - for APP in build_ios_${{ matrix.platform }}/bin/*_test.app; do + for test_name in diskann_mobile_compat_test diskann_mobile_collection_test; do + test_app="$BUILD_DIR/bin/${test_name}.app" + if [ ! -d "$test_app" ]; then + echo "::error::Missing required iOS DiskAnn test app: $test_app" + exit 1 + fi + done + + for APP in "$BUILD_DIR"/bin/*_test.app; do [ -d "$APP" ] || continue TEST_NAME=$(basename "$APP" .app) + + # All test cases in this legacy target are disabled with #if 0. + # Do not count an empty GTest binary as an executed test suite. + if [ "$TEST_NAME" = "cosine_distance_matrix_int8_test" ]; then + echo "::notice::Skipping ${TEST_NAME}: it contains no enabled GTest cases" + continue + fi + BUNDLE_ID="com.zvec.${TEST_NAME}" + LOG_FILE="$RUNNER_TEMP/${TEST_NAME}.log" TOTAL=$((TOTAL + 1)) echo "::group::Running ${TEST_NAME}" xcrun simctl install "$DEVICE_ID" "$APP" - set +eo pipefail for attempt in 1 2 3; do - xcrun simctl launch --console "$DEVICE_ID" "$BUNDLE_ID" 2>&1 | tee /tmp/${TEST_NAME}.log - LAUNCH_EXIT=${PIPESTATUS[0]} - if ! grep -q "unknown to FrontBoard" /tmp/${TEST_NAME}.log; then + set +e + python3 - "$DEVICE_ID" "$BUNDLE_ID" "$LOG_FILE" <<'PY' + import os + import pathlib + import subprocess + import sys + + device_id, bundle_id, log_file = sys.argv[1:] + env = os.environ.copy() + env["SIMCTL_CHILD_GTEST_COLOR"] = "no" + try: + result = subprocess.run( + ["xcrun", "simctl", "launch", "--console", device_id, bundle_id], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=600, + ) + output = result.stdout or "" + return_code = result.returncode + except subprocess.TimeoutExpired as error: + output = error.stdout or "" + if isinstance(output, bytes): + output = output.decode("utf-8", errors="replace") + output += "\nTest timed out after 600 seconds.\n" + return_code = 124 + + pathlib.Path(log_file).write_text(output, encoding="utf-8") + print(output, end="") + sys.exit(return_code) + PY + LAUNCH_EXIT=$? + set -e + + SUMMARY_FOUND=0 + if grep -Eq '^\[==========\][[:space:]]+[0-9]+ tests?[[:space:]]+from[[:space:]]+[0-9]+ test (suites?|cases?)[[:space:]]+ran\.' "$LOG_FILE"; then + SUMMARY_FOUND=1 + elif [ "$TEST_NAME" = "c_api_test" ] && grep -Eq '^Failed: [0-9]+[[:space:]]*$' "$LOG_FILE"; then + SUMMARY_FOUND=1 + fi + + RETRY_REASON="" + if grep -q "unknown to FrontBoard" "$LOG_FILE"; then + RETRY_REASON="FrontBoard has not registered ${TEST_NAME} yet" + elif [ "$LAUNCH_EXIT" -eq 0 ] && [ "$SUMMARY_FOUND" -eq 0 ] && \ + ! grep -Eq '^\[ FAILED \]' "$LOG_FILE"; then + RETRY_REASON="${TEST_NAME} exited without a complete test summary" + else break fi - echo "::warning::Attempt ${attempt}/3: FrontBoard has not registered ${TEST_NAME} yet, retrying in 3s..." + + if [ "$attempt" -eq 3 ]; then + break + fi + + echo "::warning::Attempt ${attempt}/3: ${RETRY_REASON}; reinstalling and retrying in 3s..." + xcrun simctl terminate "$DEVICE_ID" "$BUNDLE_ID" 2>/dev/null || true + xcrun simctl uninstall "$DEVICE_ID" "$BUNDLE_ID" 2>/dev/null || true sleep 3 + xcrun simctl install "$DEVICE_ID" "$APP" done - set -eo pipefail - if grep -q '\[ FAILED \]' /tmp/${TEST_NAME}.log; then + EXPECTED_COUNT="" + case "$TEST_NAME" in + diskann_mobile_compat_test) + EXPECTED_COUNT=10 + ;; + diskann_mobile_collection_test) + EXPECTED_COUNT=4 + ;; + esac + + if [ "$LAUNCH_EXIT" -eq 124 ]; then + echo "::error::${TEST_NAME} timed out after 600 seconds" + xcrun simctl terminate "$DEVICE_ID" "$BUNDLE_ID" 2>/dev/null || true + FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}" + elif [ "$LAUNCH_EXIT" -ne 0 ]; then + echo "::error::${TEST_NAME} launch exited ${LAUNCH_EXIT}" + FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}" + elif grep -Eq '^\[ FAILED \]' "$LOG_FILE"; then echo "::error::${TEST_NAME} has failing tests" FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}" - elif grep -q '\[ PASSED \]' /tmp/${TEST_NAME}.log; then + elif [ -n "$EXPECTED_COUNT" ] && \ + ! grep -Eq "^\\[ PASSED \\][[:space:]]+${EXPECTED_COUNT} tests?\\.[[:space:]]*$" "$LOG_FILE"; then + echo "::error::${TEST_NAME} did not report the required ${EXPECTED_COUNT}/${EXPECTED_COUNT} passing tests" + FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}" + elif [ -n "$EXPECTED_COUNT" ] && \ + ! grep -Eq "^\\[==========\\][[:space:]]+${EXPECTED_COUNT} tests?[[:space:]]+from[[:space:]]+[1-9][0-9]* test (suites?|cases?)[[:space:]]+ran\\." "$LOG_FILE"; then + echo "::error::${TEST_NAME} did not run exactly ${EXPECTED_COUNT} tests" + FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}" + elif grep -Eq '^\[==========\][[:space:]]+[1-9][0-9]* tests?[[:space:]]+from[[:space:]]+[1-9][0-9]* test (suites?|cases?)[[:space:]]+ran\.' "$LOG_FILE"; then PASSED=$((PASSED + 1)) - elif grep -qE 'Failed: 0$' /tmp/${TEST_NAME}.log; then + if [ "$TEST_NAME" = "diskann_mobile_compat_test" ]; then + DISKANN_COMPAT_PASSED=1 + elif [ "$TEST_NAME" = "diskann_mobile_collection_test" ]; then + DISKANN_COLLECTION_PASSED=1 + fi + elif [ "$TEST_NAME" = "c_api_test" ] && \ + grep -Eq '^Passed: [1-9][0-9]*[[:space:]]*$' "$LOG_FILE" && \ + grep -Eq '^Failed: 0[[:space:]]*$' "$LOG_FILE"; then # c_api_test uses a custom test framework (not GTest) PASSED=$((PASSED + 1)) - elif [ "$LAUNCH_EXIT" -eq 0 ]; then - echo "::warning::${TEST_NAME} exited 0 but produced no recognisable test summary" - PASSED=$((PASSED + 1)) else - echo "::error::${TEST_NAME} exited ${LAUNCH_EXIT} with no test summary" + echo "::error::${TEST_NAME} produced no recognisable passing test summary" FAILED_TESTS="${FAILED_TESTS} ${TEST_NAME}" fi echo "::endgroup::" done echo "Test summary: ${PASSED}/${TOTAL} passed" + if [ "$TOTAL" -eq 0 ]; then + echo "::error::No iOS test apps were discovered" + exit 1 + fi if [ -n "$FAILED_TESTS" ]; then echo "::error::Failed tests:${FAILED_TESTS}" exit 1 fi + if [ "$DISKANN_COMPAT_PASSED" -ne 1 ] || [ "$DISKANN_COLLECTION_PASSED" -ne 1 ]; then + echo "::error::Required DiskAnn tests did not both complete successfully" + exit 1 + fi + if [ "$PASSED" -ne "$TOTAL" ]; then + echo "::error::Only ${PASSED}/${TOTAL} iOS tests completed successfully" + exit 1 + fi - name: Shutdown Simulator if: matrix.test_on_simulator && always() run: | - xcrun simctl shutdown "$DEVICE_ID" || true + if [ -n "${DEVICE_ID:-}" ]; then + xcrun simctl shutdown "$DEVICE_ID" || true + fi diff --git a/.github/workflows/_build_prebuilt.yml b/.github/workflows/_build_prebuilt.yml index 74705f9b5..77b10a3b5 100644 --- a/.github/workflows/_build_prebuilt.yml +++ b/.github/workflows/_build_prebuilt.yml @@ -397,14 +397,15 @@ jobs: run: | NDK="${ANDROID_NDK_LATEST_HOME:-$ANDROID_NDK_HOME}" echo "Using NDK: $NDK" - # c++_static matches scripts/build_android.sh since #627: the - # slimmed libraries are self-contained, no shared libc++ needed. + # The packaged C++ API exchanges STL objects across the shared-library + # boundary, so zvec and its consumers must use the same libc++. + # c++_static is reserved for the build-tree static SDK targets. cmake -S . -B build -G Ninja \ -DCMAKE_TOOLCHAIN_FILE="$NDK/build/cmake/android.toolchain.cmake" \ -DANDROID_NDK="$NDK" \ -DANDROID_ABI="$ZVEC_ANDROID_ABI" \ -DANDROID_NATIVE_API_LEVEL="$ZVEC_ANDROID_API_LEVEL" \ - -DANDROID_STL=c++_static \ + -DANDROID_STL=c++_shared \ -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_INSTALL_PREFIX="$ZVEC_INSTALL_DIR" \ -DBUILD_TOOLS=OFF \ @@ -437,8 +438,28 @@ jobs: run: | mkdir -p "$ZVEC_DIST_DIR" cp -R "$ZVEC_INSTALL_DIR"/. "$ZVEC_DIST_DIR"/. + NDK="${ANDROID_NDK_LATEST_HOME:-$ANDROID_NDK_HOME}" + case "$ZVEC_ANDROID_ABI" in + arm64-v8a) LIBCXX_TRIPLE=aarch64-linux-android ;; + armeabi-v7a) LIBCXX_TRIPLE=arm-linux-androideabi ;; + x86) LIBCXX_TRIPLE=i686-linux-android ;; + x86_64) LIBCXX_TRIPLE=x86_64-linux-android ;; + *) + echo "Unsupported Android ABI: $ZVEC_ANDROID_ABI" >&2 + exit 1 + ;; + esac + LIBCXX_SHARED=$(find "$NDK/toolchains/llvm/prebuilt" -type f \ + -path "*/sysroot/usr/lib/$LIBCXX_TRIPLE/libc++_shared.so" \ + -print -quit) + if [ -z "$LIBCXX_SHARED" ]; then + echo "Unable to locate libc++_shared.so for $ZVEC_ANDROID_ABI" >&2 + exit 1 + fi + cp "$LIBCXX_SHARED" "$ZVEC_DIST_DIR/lib/libc++_shared.so" # Ship shared libraries only (same policy as desktop platforms). rm -f "$ZVEC_DIST_DIR"/lib/*.a + test -f "$ZVEC_DIST_DIR/lib/libc++_shared.so" - name: Smoke-test staged SDK (link only) shell: bash @@ -451,6 +472,7 @@ jobs: -DCMAKE_TOOLCHAIN_FILE="$NDK/build/cmake/android.toolchain.cmake" \ -DANDROID_ABI="$ZVEC_ANDROID_ABI" \ -DANDROID_NATIVE_API_LEVEL="$ZVEC_ANDROID_API_LEVEL" \ + -DANDROID_STL=c++_shared \ -DCMAKE_BUILD_TYPE=Release \ -DZVEC_SDK_DIR="$ZVEC_DIST_DIR" cmake --build build_smoke --parallel 2 diff --git a/CMakeLists.txt b/CMakeLists.txt index 07ec8ac96..da72f7c50 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -110,13 +110,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}") @@ -124,6 +133,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}") if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64|AMD64" AND NOT ANDROID AND NOT IOS) include(CheckCXXCompilerFlag) @@ -159,15 +169,48 @@ message(STATUS "RABITQ_SUPPORTED: ${RABITQ_SUPPORTED}") # DiskAnn support: # - Linux x86_64 and ARM64 with io_uring, libaio, or pread +# - 64-bit Android and iOS with the portable synchronous pread backend # - macOS ARM64 (Apple Silicon) with synchronous pread -if((CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64|AMD64|aarch64|arm64)$" AND NOT ANDROID AND NOT IOS) - OR (CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(arm64|aarch64)$" AND NOT IOS)) +# - Windows x86_64 with overlapped I/O +set(DISKANN_SUPPORTED OFF) +set(_diskann_unsupported_reason + "unsupported target ${CMAKE_SYSTEM_NAME}/${CMAKE_SYSTEM_PROCESSOR}") + +# All supported targets require 64-bit pointers. This is also the authoritative +# gate for Android, whose ABI-specific processor names vary. +if(NOT CMAKE_SIZEOF_VOID_P EQUAL 8) + set(_diskann_unsupported_reason + "32-bit targets are not supported (CMAKE_SIZEOF_VOID_P=${CMAKE_SIZEOF_VOID_P})") +elseif(ANDROID OR IOS) set(DISKANN_SUPPORTED ON) +elseif(CMAKE_SYSTEM_NAME STREQUAL "Linux") + if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64|AMD64|aarch64|arm64)$") + set(DISKANN_SUPPORTED ON) + else() + set(_diskann_unsupported_reason + "Linux requires x86_64 or ARM64 (detected ${CMAKE_SYSTEM_PROCESSOR})") + endif() +elseif(CMAKE_SYSTEM_NAME STREQUAL "Darwin") + if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(arm64|aarch64)$") + set(DISKANN_SUPPORTED ON) + else() + set(_diskann_unsupported_reason + "macOS requires ARM64 (detected ${CMAKE_SYSTEM_PROCESSOR})") + endif() +elseif(WIN32) + if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64|AMD64)$") + set(DISKANN_SUPPORTED ON) + else() + set(_diskann_unsupported_reason + "Windows requires x86_64 (detected ${CMAKE_SYSTEM_PROCESSOR})") + endif() +endif() + +if(DISKANN_SUPPORTED) add_definitions(-DDISKANN_SUPPORTED=1) else() - set(DISKANN_SUPPORTED OFF) add_definitions(-DDISKANN_SUPPORTED=0) - message(STATUS "DiskAnn support disabled for ${CMAKE_SYSTEM_NAME}/${CMAKE_SYSTEM_PROCESSOR}") + message(STATUS "DiskAnn support disabled: ${_diskann_unsupported_reason}") endif() message(STATUS "DISKANN_SUPPORTED: ${DISKANN_SUPPORTED}") @@ -180,6 +223,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 ARROW.BUILD glog gflags Lz4.BUILD) diff --git a/cmake/bazel.cmake b/cmake/bazel.cmake index e6e29f189..e5b5036c4 100644 --- a/cmake/bazel.cmake +++ b/cmake/bazel.cmake @@ -616,16 +616,13 @@ 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}) if(IOS) - # iOS has no shared libraries, so the main target is static as well. - # Building a second, identical archive under the ${_NAME}_static name is - # not just wasteful, it breaks the build: giving both the same OUTPUT_NAME - # makes Ninja fail ("multiple rules generate ..."), while distinct names - # make targets that link both variants fail with duplicate symbols. - # A single archive exposed under both names avoids both problems. + # iOS has no shared libraries. Expose the single static archive under both + # target names to avoid duplicate objects while preserving callers of the + # explicit _static target. add_library( ${_NAME} STATIC ${_OPTION} $ ) @@ -634,9 +631,23 @@ macro(_add_library _NAME _OPTION) add_library( ${_NAME}_static STATIC ${_OPTION} $ ) - add_library( - ${_NAME} SHARED ${_OPTION} $ - ) + if(ANDROID AND ANDROID_STL STREQUAL "c++_static") + add_library( + ${_NAME} STATIC ${_OPTION} $ + ) + # Keep both Android static build targets but give the main archive a + # distinct file name. Canonicalization ensures only the main archive is + # whole-archived when both target names appear in a dependency graph. + 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} $ + ) + endif() add_dependencies(${_NAME} ${_NAME}_static) if(NOT MSVC) set_property(TARGET ${_NAME}_static PROPERTY OUTPUT_NAME ${_NAME}) @@ -690,6 +701,25 @@ endfunction() ## Link libraries function(_target_link_libraries _NAME) + function(_resolve_link_target LIB RESULT_VAR) + set(RESOLVED_LINK_TARGET ${LIB}) + if(TARGET ${RESOLVED_LINK_TARGET}) + get_target_property( + ALIASED_LINK_TARGET ${RESOLVED_LINK_TARGET} ALIASED_TARGET) + if(ALIASED_LINK_TARGET) + set(RESOLVED_LINK_TARGET ${ALIASED_LINK_TARGET}) + endif() + get_target_property( + CANONICAL_LINK_TARGET + ${RESOLVED_LINK_TARGET} + ZVEC_CANONICAL_LINK_TARGET) + if(CANONICAL_LINK_TARGET) + set(RESOLVED_LINK_TARGET ${CANONICAL_LINK_TARGET}) + endif() + endif() + set(${RESULT_VAR} ${RESOLVED_LINK_TARGET} 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) @@ -697,6 +727,7 @@ function(_target_link_libraries _NAME) set(LOCAL_RESULT "") foreach(LIB ${LIB_LIST}) + _resolve_link_target(${LIB} LIB) if(NOT TARGET ${LIB}) continue() endif() @@ -726,7 +757,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() @@ -736,11 +770,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 bde2d0ab8..ecb138c17 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) @@ -117,6 +141,8 @@ add_executable(diskann-core-example core/diskann_main.cc) target_link_libraries(diskann-core-example PRIVATE zvec-core-lib) list(APPEND ZVEC_EXAMPLE_TARGETS diskann-core-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 d93624917..3cd47b4c8 100644 --- a/examples/c/diskann_example.c +++ b/examples/c/diskann_example.c @@ -21,7 +21,9 @@ * a Vamana graph structure combined with product quantization (PQ) to * achieve high recall with efficient disk I/O. * - * NOTE: DiskANN is available on Linux x86_64/ARM64 and macOS ARM64. + * NOTE: DiskANN is available on Linux x86/ARM64 and macOS ARM64 (using + * io_uring, libaio, or pread), and on Android and iOS via 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/python/tests/detail/fixture_helper.py b/python/tests/detail/fixture_helper.py index 7b3f743ae..a5c5804a1 100644 --- a/python/tests/detail/fixture_helper.py +++ b/python/tests/detail/fixture_helper.py @@ -3,9 +3,13 @@ import platform DISKANN_SUPPORTED = ( - platform.system() == "Linux" - and platform.machine() in ("x86_64", "AMD64", "aarch64", "arm64") -) or (platform.system() == "Darwin" and platform.machine() in ("aarch64", "arm64")) + ( + platform.system() == "Linux" + and platform.machine() in ("x86_64", "AMD64", "aarch64", "arm64") + ) + or (platform.system() == "Darwin" and platform.machine() in ("aarch64", "arm64")) + or (platform.system() == "Windows" and platform.machine() in ("x86_64", "AMD64")) +) from typing import Any, Generator from zvec.typing import DataType, StatusCode, MetricType, QuantizeType @@ -26,7 +30,8 @@ def _ensure_diskann_runtime_or_reason() -> str | None: if not DISKANN_SUPPORTED: _DISKANN_PRELOAD_REASON = ( - "DiskAnn is supported on Linux (x86_64/ARM64) and macOS ARM64" + "DiskAnn is supported on Linux (x86_64/ARM64), macOS ARM64, " + "and Windows x86_64" ) return _DISKANN_PRELOAD_REASON _DISKANN_PRELOAD_REASON = None diff --git a/python/tests/test_collection_diskann.py b/python/tests/test_collection_diskann.py index 7ad583130..843dbf756 100644 --- a/python/tests/test_collection_diskann.py +++ b/python/tests/test_collection_diskann.py @@ -15,11 +15,9 @@ Mirrors ``test_collection_hnsw_rabitq.py`` but targets the DiskAnn index. -DiskAnn must be built for Linux (x86_64/ARM64) or macOS ARM64. Other -platforms are skipped wholesale. - -If the prerequisite fails the whole module is skipped so the rest of the -test suite is not affected. macOS uses synchronous pread. +DiskAnn must be built for Linux (x86_64/ARM64), macOS ARM64, or Windows +x86_64. Other platforms are skipped wholesale. Linux selects io_uring, +libaio, or pread; macOS uses pread; Windows uses overlapped I/O. """ from __future__ import annotations @@ -40,8 +38,11 @@ and platform.machine() in ("x86_64", "AMD64", "aarch64", "arm64") ) or (sys.platform == "darwin" and platform.machine() in ("aarch64", "arm64")) + or (sys.platform == "win32" and platform.machine() in ("x86_64", "AMD64")) + ), + reason=( + "DiskAnn is supported on Linux (x86_64/ARM64), macOS ARM64, and Windows x86_64" ), - reason="DiskAnn is supported on Linux (x86_64/ARM64) and macOS ARM64", ) import zvec # noqa: E402 diff --git a/python/tests/test_typing.py b/python/tests/test_typing.py index f6a3f3788..4b79da17b 100644 --- a/python/tests/test_typing.py +++ b/python/tests/test_typing.py @@ -38,6 +38,8 @@ (IndexType.HNSW, "HNSW"), (IOBackendType.PREAD, "PREAD"), (IOBackendType.IO_URING, "IO_URING"), + (IOBackendType.WINDOWS_OVERLAPPED, "WINDOWS_OVERLAPPED"), + (IOBackendType.UNAVAILABLE, "UNAVAILABLE"), (MetricType.COSINE, "COSINE"), (QuantizeType.INT8, "INT8"), (StatusCode.OK, "OK"), @@ -54,6 +56,8 @@ def test_enum_names(member, name): (IndexType.HNSW, 1), (IOBackendType.PREAD, 0), (IOBackendType.IO_URING, 2), + (IOBackendType.WINDOWS_OVERLAPPED, 3), + (IOBackendType.UNAVAILABLE, 4), (MetricType.COSINE, 3), (QuantizeType.INT8, 2), (StatusCode.OK, 0), @@ -118,7 +122,10 @@ def test_index_type_has_member(member): assert member in IndexType.__members__ -@pytest.mark.parametrize("member", ["PREAD", "LIBAIO", "IO_URING"]) +@pytest.mark.parametrize( + "member", + ["PREAD", "LIBAIO", "IO_URING", "WINDOWS_OVERLAPPED", "UNAVAILABLE"], +) def test_io_backend_type_has_member(member): assert member in IOBackendType.__members__ @@ -127,7 +134,10 @@ def test_current_io_backend_type(): backend = zvec.io_backend_type() assert isinstance(backend, IOBackendType) assert zvec.io_backend_description() - if platform.system() == "Darwin": + if platform.system() == "Windows": + assert backend == IOBackendType.WINDOWS_OVERLAPPED + assert "overlapped" in zvec.io_backend_description().lower() + elif platform.system() == "Darwin": assert backend == IOBackendType.PREAD assert "pread" in zvec.io_backend_description().lower() diff --git a/python/zvec/__init__.pyi b/python/zvec/__init__.pyi index 507ff47dc..c9dff99de 100644 --- a/python/zvec/__init__.pyi +++ b/python/zvec/__init__.pyi @@ -56,6 +56,8 @@ def io_backend_type() -> IOBackendType: Linux selects IOBackendType.IO_URING, IOBackendType.LIBAIO, or IOBackendType.PREAD in that order. macOS ARM64 uses IOBackendType.PREAD. + Windows x86_64 uses IOBackendType.WINDOWS_OVERLAPPED. Unsupported target + architectures return IOBackendType.UNAVAILABLE. """ def io_backend_description() -> str: @@ -63,7 +65,8 @@ def io_backend_description() -> str: The description identifies io_uring, libaio, or pread. On Linux, the pread description includes guidance for enabling io_uring or installing - libaio. + libaio. Windows reports its overlapped-I/O backend. Unsupported target + architectures report that DiskAnn is unavailable. """ def set_default_jieba_dict_dir(dir: str) -> None: diff --git a/python/zvec/typing/__init__.pyi b/python/zvec/typing/__init__.pyi index eaa7326cb..a2f735995 100644 --- a/python/zvec/typing/__init__.pyi +++ b/python/zvec/typing/__init__.pyi @@ -131,6 +131,9 @@ class IOBackendType: - PREAD: Synchronous pread() — no async I/O. - LIBAIO: libaio loaded at runtime via dlopen(). - IO_URING: io_uring via raw kernel syscalls (zero dependency). + - WINDOWS_OVERLAPPED: Windows unbuffered overlapped I/O using per-context + I/O completion ports. + - UNAVAILABLE: DiskAnn is disabled on this target architecture. Examples: >>> from zvec.typing import IOBackendType @@ -145,14 +148,24 @@ class IOBackendType: LIBAIO IO_URING + + WINDOWS_OVERLAPPED + + UNAVAILABLE """ - IO_URING: typing.ClassVar[IOBackendType] # value = - LIBAIO: typing.ClassVar[IOBackendType] # value = PREAD: typing.ClassVar[IOBackendType] # value = + LIBAIO: typing.ClassVar[IOBackendType] # value = + IO_URING: typing.ClassVar[IOBackendType] # value = + WINDOWS_OVERLAPPED: typing.ClassVar[ + IOBackendType + ] # value = + UNAVAILABLE: typing.ClassVar[ + IOBackendType + ] # value = __members__: typing.ClassVar[ dict[str, IOBackendType] - ] # value = {'PREAD': , 'LIBAIO': , 'IO_URING': } + ] # value includes PREAD, LIBAIO, IO_URING, WINDOWS_OVERLAPPED, and UNAVAILABLE def __eq__(self, other: typing.Any) -> bool: ... def __getstate__(self) -> int: ... diff --git a/scripts/build_android.sh b/scripts/build_android.sh index 54337f92f..a7de8ccd5 100755 --- a/scripts/build_android.sh +++ b/scripts/build_android.sh @@ -58,8 +58,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" \ @@ -360,4 +364,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 0f28f7431..51cbd75a1 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -215,6 +215,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/ailego/io/io_backend_def.h b/src/ailego/io/io_backend_def.h index d5636e5a7..87eb52d5b 100644 --- a/src/ailego/io/io_backend_def.h +++ b/src/ailego/io/io_backend_def.h @@ -48,6 +48,10 @@ inline const char *IOBackendTypeName(IOBackendType type) { return "libaio"; case IOBackendType::kPread: return "pread"; + case IOBackendType::kWindowsOverlapped: + return "windows_overlapped"; + case IOBackendType::kUnavailable: + return "unavailable"; } return "unknown"; } @@ -71,6 +75,11 @@ inline const char *IOBackendDescription(IOBackendType type) { #else return "Synchronous pread() I/O backend."; #endif + case IOBackendType::kWindowsOverlapped: + return "windows_overlapped: Windows unbuffered overlapped I/O backend " + "using per-context I/O completion ports."; + case IOBackendType::kUnavailable: + return "unavailable: DiskAnn is disabled on this target."; } return "Unknown I/O backend."; } @@ -88,11 +97,20 @@ class IOBackend { } // Returns the active backend, probing on the first call. Linux prefers - // io_uring, then libaio, then pread; macOS ARM64 uses pread. + // io_uring, then libaio, then pread; Windows uses overlapped I/O; macOS + // ARM64, Android and iOS use pread. + // + // Android is deliberately excluded from async probing: its seccomp sandbox + // may not permit the io_uring_setup() syscall, and a blocked syscall raises + // SIGSYS rather than returning an error, so probing could crash the process. IOBackendType available() { std::call_once(probe_once_, [this]() { IOBackendType selected = IOBackendType::kPread; -#if defined(__linux) || defined(__linux__) +#if !defined(DISKANN_SUPPORTED) || !DISKANN_SUPPORTED + selected = IOBackendType::kUnavailable; +#elif defined(_WIN32) || defined(_WIN64) + selected = IOBackendType::kWindowsOverlapped; +#elif (defined(__linux) || defined(__linux__)) && !defined(__ANDROID__) if (io_uring_supported()) { selected = IOBackendType::kIoUring; } else if (LibAioLoader::Instance().load() && @@ -117,6 +135,20 @@ class IOBackend { return available() == IOBackendType::kIoUring; } + // Persist a per-context setup fallback as the process-wide selection. The + // Linux enum values are ordered from the synchronous backend to the most + // capable asynchronous backend, so this operation is monotonic: a racing + // successful setup can never promote the process after another context has + // demonstrated that the preferred backend is unavailable at runtime. + void downgrade(IOBackendType fallback) { + IOBackendType current = type_.load(std::memory_order_acquire); + while (static_cast(fallback) < static_cast(current) && + !type_.compare_exchange_weak(current, fallback, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + } + } + // Returns the cached backend type without triggering the probe. IOBackendType type() const { return type_.load(std::memory_order_acquire); @@ -135,7 +167,7 @@ class IOBackend { private: IOBackend() = default; -#if defined(__linux) || defined(__linux__) +#if (defined(__linux) || defined(__linux__)) && !defined(__ANDROID__) // Probe io_uring availability with a minimal ring setup using only raw // syscalls — no dependency on liburing. A successful setup alone is NOT // sufficient: io_uring_setup() exists since Linux 5.1, but the read path diff --git a/src/binding/c/c_api.cc b/src/binding/c/c_api.cc index 6281cd295..9ca59ce60 100644 --- a/src/binding/c/c_api.cc +++ b/src/binding/c/c_api.cc @@ -766,6 +766,23 @@ const char *zvec_get_default_jieba_dict_dir(void) { // I/O Backend Introspection // ============================================================================= +static_assert( + static_cast(zvec::ailego::IOBackendType::kPread) == + ZVEC_IO_BACKEND_TYPE_PREAD); +static_assert( + static_cast(zvec::ailego::IOBackendType::kLibAio) == + ZVEC_IO_BACKEND_TYPE_LIBAIO); +static_assert( + static_cast(zvec::ailego::IOBackendType::kIoUring) == + ZVEC_IO_BACKEND_TYPE_IO_URING); +static_assert( + static_cast( + zvec::ailego::IOBackendType::kWindowsOverlapped) == + ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED); +static_assert( + static_cast(zvec::ailego::IOBackendType::kUnavailable) == + ZVEC_IO_BACKEND_TYPE_UNAVAILABLE); + zvec_io_backend_type_t zvec_get_io_backend_type(void) { auto type = zvec::ailego::current_io_backend_type(); return static_cast(static_cast(type)); diff --git a/src/binding/python/CMakeLists.txt b/src/binding/python/CMakeLists.txt index c05641e72..39b437726 100644 --- a/src/binding/python/CMakeLists.txt +++ b/src/binding/python/CMakeLists.txt @@ -46,22 +46,15 @@ endif() if (CMAKE_SYSTEM_NAME STREQUAL "Linux") target_link_libraries(_zvec PRIVATE -Wl,--whole-archive - $ - $ - $ - $ - $ - $ - $ - $ - $ - $ - $ - $ - $ - $ + $ + $ + $ + $ -Wl,--no-whole-archive zvec + zvec_core + zvec_ailego + zvec_turbo ${CMAKE_DL_LIBS} ) target_link_options(_zvec PRIVATE @@ -84,43 +77,27 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Linux") # DiskAnn runtime library. elseif (APPLE) target_link_libraries(_zvec PRIVATE - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ - -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ + -Wl,-force_load,$ zvec + zvec_core + zvec_ailego + zvec_turbo ) target_link_libraries(_zvec PRIVATE -Wl,-exported_symbols_list,${CMAKE_CURRENT_SOURCE_DIR}/exports.mac ) elseif (MSVC) set(_zvec_whole_archive_libs - core_knn_flat_static - core_knn_flat_sparse_static - core_knn_hnsw_static - core_knn_hnsw_sparse_static - core_knn_ivf_static - core_knn_vamana_static - core_knn_cluster_static - core_knn_diskann_static - core_mix_reducer_static - core_metric_static - core_utility_static - core_quantizer_static + zvec + zvec_core + zvec_ailego + zvec_turbo ) target_link_libraries(_zvec PRIVATE ${_zvec_whole_archive_libs} - zvec ) foreach(_lib ${_zvec_whole_archive_libs}) target_link_options(_zvec PRIVATE diff --git a/src/binding/python/model/common/python_config.cc b/src/binding/python/model/common/python_config.cc index 4973167c0..3e5e7e58a 100644 --- a/src/binding/python/model/common/python_config.cc +++ b/src/binding/python/model/common/python_config.cc @@ -220,7 +220,8 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { "Read the currently registered default jieba dict directory."); // Returns the selected DiskAnn I/O backend. Linux tries io_uring, then - // libaio, then pread; macOS ARM64 uses pread. + // libaio, then pread; macOS ARM64 uses pread. Unsupported targets report + // that DiskAnn is unavailable. m.def( "io_backend_type", []() -> ailego::IOBackendType { @@ -230,7 +231,9 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { "as an IOBackendType enum (zvec.typing.IOBackendType). " "Linux selects IOBackendType.IO_URING, IOBackendType.LIBAIO, or " "IOBackendType.PREAD in that order. macOS ARM64 uses " - "IOBackendType.PREAD."); + "IOBackendType.PREAD. Windows x86_64 uses " + "IOBackendType.WINDOWS_OVERLAPPED. Unsupported target architectures " + "return IOBackendType.UNAVAILABLE."); // Returns a human-readable description identifying io_uring, libaio, or // pread, with asynchronous-backend guidance for Linux pread fallback. @@ -240,7 +243,9 @@ void ZVecPyConfig::Initialize(pybind11::module_ &m) { "Returns a human-readable description of the current I/O backend. " "The description identifies io_uring, libaio, or pread. On Linux, the " "pread description includes guidance for enabling io_uring or " - "installing libaio."); + "installing libaio. Windows reports its overlapped-I/O backend. " + "Unsupported target architectures report that DiskAnn is " + "unavailable."); } diff --git a/src/binding/python/model/param/python_param.cc b/src/binding/python/model/param/python_param.cc index 35c2b8a09..5c3919301 100644 --- a/src/binding/python/model/param/python_param.cc +++ b/src/binding/python/model/param/python_param.cc @@ -1071,7 +1071,8 @@ and accuracy. use_soar (bool): Whether to enable SOAR (Scalable Optimized Adaptive Routing) for improved IVF search performance. Default is False. quantize_type (QuantizeType): Optional quantization type for vector - compression (e.g., FP16, INT8). Default is ``QuantizeType.UNDEFINED``. + compression. DiskAnn currently supports ``QuantizeType.FP16`` only. + Default is ``QuantizeType.UNDEFINED``. Examples: >>> from zvec.typing import MetricType, QuantizeType @@ -1107,8 +1108,9 @@ Constructs an IVFIndexParam instance. n_iters (int, optional): Number of k-means iterations during training. Defaults to 10. use_soar (bool, optional): Enable SOAR optimization. Defaults to False. - quantize_type (QuantizeType, optional): Vector quantization type. - Defaults to QuantizeType.UNDEFINED. + quantize_type (QuantizeType, optional): Vector quantization type. DiskAnn + currently supports QuantizeType.FP16 only. Defaults to + QuantizeType.UNDEFINED. quantizer_param (QuantizerParam, optional): Quantizer configuration. Defaults to QuantizerParam(). )pbdoc") diff --git a/src/binding/python/typing/python_type.cc b/src/binding/python/typing/python_type.cc index 971e2059c..b48b31bd1 100644 --- a/src/binding/python/typing/python_type.cc +++ b/src/binding/python/typing/python_type.cc @@ -148,6 +148,9 @@ Enumeration of supported I/O backend types for DiskAnn disk reads. - PREAD: Synchronous pread(); no async I/O. - LIBAIO: libaio loaded at runtime via dlopen(). - IO_URING: io_uring via raw kernel syscalls (zero dependency). +- WINDOWS_OVERLAPPED: Windows unbuffered overlapped I/O using per-context + I/O completion ports. +- UNAVAILABLE: DiskAnn is disabled on this target architecture. Examples: >>> from zvec.typing import IOBackendType @@ -156,7 +159,9 @@ Enumeration of supported I/O backend types for DiskAnn disk reads. )pbdoc") .value("PREAD", ailego::IOBackendType::kPread) .value("LIBAIO", ailego::IOBackendType::kLibAio) - .value("IO_URING", ailego::IOBackendType::kIoUring); + .value("IO_URING", ailego::IOBackendType::kIoUring) + .value("WINDOWS_OVERLAPPED", ailego::IOBackendType::kWindowsOverlapped) + .value("UNAVAILABLE", ailego::IOBackendType::kUnavailable); } void ZVecPyTyping::bind_status(py::module_ &m) { diff --git a/src/core/algorithm/CMakeLists.txt b/src/core/algorithm/CMakeLists.txt index 7931bf344..22050f9e9 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 supports Linux (x86_64/ARM64) and macOS ARM64\n" + "// DiskAnn supports Linux x86_64/ARM64, macOS ARM64, 64-bit Android/iOS, and Windows x86_64\n" "namespace zvec { namespace core { /* empty namespace for compatibility */ } }\n" ) diff --git a/src/core/algorithm/diskann/CMakeLists.txt b/src/core/algorithm/diskann/CMakeLists.txt index 9dcab7da0..0ff7f72d4 100644 --- a/src/core/algorithm/diskann/CMakeLists.txt +++ b/src/core/algorithm/diskann/CMakeLists.txt @@ -13,14 +13,15 @@ file(GLOB_RECURSE ALL_SRCS *.cc *.c) # loaded at runtime via dlopen()/dlsym() (see libaio_loader.h), so we do NOT # link against -laio. On Linux, ${CMAKE_DL_LIBS} provides the system-level # dependency for dlopen/dlsym/dlclose. macOS ARM64 uses pread and needs no -# extra I/O backend library. -set(CORE_KNN_DISKANN_LIBS core_framework core_knn_cluster) +# extra I/O backend library. Windows overlapped I/O is built in and likewise +# needs no extra backend library. +set(CORE_KNN_DISKANN_LIBS zvec_ailego core_framework core_knn_cluster) if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64|AMD64|aarch64|arm64)$") list(APPEND CORE_KNN_DISKANN_LIBS ${CMAKE_DL_LIBS}) endif() -if(NOT APPLE) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") set(CORE_KNN_DISKANN_LDFLAGS "-Wl,--exclude-libs,libparquet.a:libarrow.a:libarrow_bundled_dependencies.a") endif() @@ -34,3 +35,11 @@ cc_library( LDFLAGS "${CORE_KNN_DISKANN_LDFLAGS}" VERSION "${PROXIMA_ZVEC_VERSION}" ) +if(MSVC) + # NOMINMAX must reach the _objects target, since _add_library() compiles + # sources there; the SHARED/STATIC targets only bundle object files and a + # PRIVATE define on them never reaches the compiler. Without this, the + # windows.h min/max macros collide with std::min/std::max and + # numeric_limits::max()/min() (error C4003). + target_compile_definitions(core_knn_diskann_objects PRIVATE NOMINMAX) +endif() diff --git a/src/core/algorithm/diskann/diskann_builder.cc b/src/core/algorithm/diskann/diskann_builder.cc index da76aee09..dc8b7e808 100644 --- a/src/core/algorithm/diskann/diskann_builder.cc +++ b/src/core/algorithm/diskann/diskann_builder.cc @@ -46,6 +46,17 @@ int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { params.get(PARAM_DISKANN_BUILDER_LIST_SIZE, &list_size_); params.get(PARAM_DISKANN_BUILDER_THREAD_COUNT, &build_thread_count_); + const double max_build_degree = + std::ceil(static_cast(max_degree_) * + static_cast(DiskAnnEntity::kDefaultGraphSlackFactor)); + if (max_degree_ == 0 || list_size_ == 0 || + max_build_degree > + static_cast((std::numeric_limits::max)() - 1U)) { + LOG_ERROR("Invalid DiskAnn graph parameters: max_degree=%u list_size=%u", + max_degree_, list_size_); + return IndexError_InvalidArgument; + } + if (build_thread_count_ == 0) { build_thread_count_ = std::max(1U, std::thread::hardware_concurrency()); } @@ -90,6 +101,12 @@ int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { if (params.has(PARAM_DISKANN_BUILDER_TRAIN_SAMPLE_RATIO)) { params.get(PARAM_DISKANN_BUILDER_TRAIN_SAMPLE_RATIO, &train_sample_ratio_); } + if (max_train_sample_count_ == 0 || !std::isfinite(train_sample_ratio_) || + train_sample_ratio_ <= 0.0 || train_sample_ratio_ > 1.0) { + LOG_ERROR("Invalid DiskAnn PQ sampling parameters: max_samples=%u ratio=%f", + max_train_sample_count_, train_sample_ratio_); + return IndexError_InvalidArgument; + } raw_meta_ = meta; @@ -141,8 +158,8 @@ int DiskAnnBuilder::init(const IndexMeta &meta, const ailego::Params ¶ms) { algo_ = DiskAnnAlgorithm::UPointer(new DiskAnnAlgorithm(entity_, max_degree_)); - trainer_ = - DiskAnnPqTrainer::UPointer(new DiskAnnPqTrainer(max_train_sample_count_)); + trainer_ = DiskAnnPqTrainer::UPointer( + new DiskAnnPqTrainer(max_train_sample_count_, train_sample_ratio_)); state_ = BUILD_STATE_INITED; @@ -304,6 +321,16 @@ int DiskAnnBuilder::calculate_pq_chunk_num() { return 0; } +bool DiskAnnBuilder::record_worker_error(int error_code) { + std::lock_guard lock(mutex_); + if (error_.load(std::memory_order_relaxed)) { + return false; + } + errcode_ = error_code != 0 ? error_code : IndexError_Runtime; + error_.store(true, std::memory_order_release); + return true; +} + int DiskAnnBuilder::build_internal(IndexThreads::Pointer threads) { auto task_group = threads->make_group(); if (!task_group) { @@ -319,24 +346,23 @@ int DiskAnnBuilder::build_internal(IndexThreads::Pointer threads) { { std::unique_lock lk(mutex_); - while (finished.load() < entity_.doc_cnt()) { + while (finished.load() < entity_.doc_cnt() && + !error_.load(std::memory_order_acquire)) { cond_.wait_until(lk, std::chrono::system_clock::now() + std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to build index while waiting finish"); - return errcode_; - } LOG_INFO("Built cnt %zu, finished percent %.3f%%", (size_t)finished.load(), finished.load() * 100.0f / entity_.doc_cnt()); } } + // Every task captures the address of the local progress counter. Never + // return while a task can still access it, including failure paths. + task_group->wait_finish(); if (error_.load(std::memory_order_acquire)) { LOG_ERROR("Failed to build index while waiting finish"); return errcode_; } - task_group->wait_finish(); return 0; } @@ -356,24 +382,21 @@ int DiskAnnBuilder::prune_internal(IndexThreads::Pointer threads) { { std::unique_lock lk(mutex_); - while (finished.load() < entity_.doc_cnt()) { + while (finished.load() < entity_.doc_cnt() && + !error_.load(std::memory_order_acquire)) { cond_.wait_until(lk, std::chrono::system_clock::now() + std::chrono::seconds(check_interval_secs_)); - if (error_.load(std::memory_order_acquire)) { - LOG_ERROR("Failed to prune index while waiting finish"); - return errcode_; - } LOG_INFO("Prune cnt %zu, finished percent %.3f%%", (size_t)finished.load(), finished.load() * 100.0f / entity_.doc_cnt()); } } + task_group->wait_finish(); if (error_.load(std::memory_order_acquire)) { LOG_ERROR("Failed to prune index while waiting finish"); return errcode_; } - task_group->wait_finish(); return 0; } @@ -433,9 +456,8 @@ void DiskAnnBuilder::do_build(uint64_t idx, size_t step_size, std::shared_ptr(&entity_, [](DiskAnnEntity *) {})); if (ailego_unlikely(ctx == nullptr)) { - if (!error_.exchange(true)) { + if (record_worker_error(IndexError_NoMemory)) { LOG_ERROR("Failed to create context"); - errcode_ = IndexError_NoMemory; } return; } @@ -444,21 +466,21 @@ void DiskAnnBuilder::do_build(uint64_t idx, size_t step_size, int ret = ctx->init(DiskAnnContext::kBuilderContext, max_degree_, pq_chunk_num_, build_meta_.element_size()); if (ailego_unlikely(ret != 0)) { - if (!error_.exchange(true)) { + if (record_worker_error(ret)) { LOG_ERROR("Failed to initialize build context"); - errcode_ = ret; } return; } ctx->set_list_size(list_size_); - for (uint64_t id = idx; id < entity_.doc_cnt(); id += step_size) { + for (uint64_t id = idx; + id < entity_.doc_cnt() && !error_.load(std::memory_order_acquire); + id += step_size) { ctx->reset_query(entity_.get_vector(id)); ret = algo_->add_node(id, ctx); if (ailego_unlikely(ret != 0)) { - if (!error_.exchange(true)) { + if (record_worker_error(ret)) { LOG_ERROR("DiskAnn graph add node failed"); - errcode_ = ret; } return; } @@ -479,9 +501,8 @@ void DiskAnnBuilder::do_prune(uint64_t idx, size_t step_size, std::shared_ptr(&entity_, [](DiskAnnEntity *) {})); if (ailego_unlikely(ctx == nullptr)) { - if (!error_.exchange(true)) { + if (record_worker_error(IndexError_NoMemory)) { LOG_ERROR("Failed to create context"); - errcode_ = IndexError_NoMemory; } return; } @@ -490,21 +511,21 @@ void DiskAnnBuilder::do_prune(uint64_t idx, size_t step_size, int ret = ctx->init(DiskAnnContext::kBuilderContext, max_degree_, pq_chunk_num_, build_meta_.element_size()); if (ailego_unlikely(ret != 0)) { - if (!error_.exchange(true)) { + if (record_worker_error(ret)) { LOG_ERROR("Failed to initialize prune context"); - errcode_ = ret; } return; } ctx->set_list_size(list_size_); - for (uint64_t id = idx; id < entity_.doc_cnt(); id += step_size) { + for (uint64_t id = idx; + id < entity_.doc_cnt() && !error_.load(std::memory_order_acquire); + id += step_size) { ctx->reset_query(entity_.get_vector(id)); ret = algo_->prune_node(id, ctx); if (ailego_unlikely(ret != 0)) { - if (!error_.exchange(true)) { + if (record_worker_error(ret)) { LOG_ERROR("DiskAnn graph add node failed"); - errcode_ = ret; } return; } @@ -540,6 +561,10 @@ int DiskAnnBuilder::train(IndexThreads::Pointer threads, LOG_ERROR("Invalid holder for DiskAnnBuilder::train"); return IndexError_InvalidArgument; } + if (!holder->is_matched(raw_meta_)) { + LOG_ERROR("Holder does not match DiskAnn builder metadata during train"); + return IndexError_Mismatch; + } LOG_INFO("Begin DiskAnnBuilder::train"); @@ -603,6 +628,10 @@ int DiskAnnBuilder::build(IndexThreads::Pointer threads, LOG_ERROR("Invalid holder for DiskAnnBuilder::build"); return IndexError_InvalidArgument; } + if (!holder->is_matched(raw_meta_)) { + LOG_ERROR("Holder does not match DiskAnn builder metadata during build"); + return IndexError_Mismatch; + } LOG_INFO("Start DiskAnnBuilder::build"); @@ -610,6 +639,28 @@ int DiskAnnBuilder::build(IndexThreads::Pointer threads, holder_ = holder; + bool build_succeeded = false; + AILEGO_DEFER([&]() { + holder_.reset(); + if (build_succeeded || entity_.doc_cnt() == 0) { + return; + } + + // A failed graph build may have appended vectors or mutated neighbors. + // Discard the partial entity and require training again before retrying. + const int reset_ret = entity_.init(raw_meta_, max_degree_, list_size_, + memory_limit_, build_thread_count_); + if (reset_ret != 0) { + LOG_ERROR("Failed to reset DiskAnn entity after build failure: %d", + reset_ret); + state_ = BUILD_STATE_INIT; + return; + } + stats_.set_trained_count(0UL); + stats_.set_trained_costtime(0UL); + state_ = BUILD_STATE_INITED; + }); + if (!threads) { threads = std::make_shared(build_thread_count_, false); @@ -624,15 +675,24 @@ int DiskAnnBuilder::build(IndexThreads::Pointer threads, return IndexError_Runtime; } - if (ailego_unlikely(holder->count() == 0)) { + const size_t declared_count = holder->count(); + if (ailego_unlikely(declared_count == 0)) { LOG_ERROR("Holder is empty"); - return IndexError_Runtime; + return IndexError_InvalidLength; } - int ret = entity_.reserve_space(holder->count()); + int ret = entity_.reserve_space(declared_count); + if (ailego_unlikely(ret != 0)) { + return ret; + } - error_ = false; + errcode_ = 0; + error_.store(false, std::memory_order_release); while (iter->is_valid()) { + if (ailego_unlikely(entity_.doc_cnt() >= declared_count)) { + LOG_ERROR("Holder contains more vectors than its declared count"); + return IndexError_InvalidLength; + } ret = entity_.add_vector(iter->key(), iter->data()); if (ailego_unlikely(ret != 0)) { return ret; @@ -640,6 +700,10 @@ int DiskAnnBuilder::build(IndexThreads::Pointer threads, iter->next(); } + if (ailego_unlikely(entity_.doc_cnt() != declared_count)) { + LOG_ERROR("Holder ended before its declared vector count"); + return IndexError_InvalidLength; + } LOG_INFO("Finished saving vector"); @@ -668,6 +732,7 @@ int DiskAnnBuilder::build(IndexThreads::Pointer threads, } state_ = BUILD_STATE_BUILT; + build_succeeded = true; stats_.set_built_count(entity_.doc_cnt()); stats_.set_built_costtime(ailego::Monotime::MilliSeconds() - start_time); @@ -694,14 +759,13 @@ int DiskAnnBuilder::dump(const IndexDumper::Pointer &dumper) { return ret; } - ret = entity_.dump(holder_, raw_meta_, dumper); + ret = entity_.dump(raw_meta_, dumper); if (ret != 0) { LOG_ERROR("Index dump failed, ret: %u", ret); - - return IndexError_Runtime; + return ret; } - stats_.set_dumped_count(holder_->count()); + stats_.set_dumped_count(entity_.doc_cnt()); stats_.set_dumped_costtime(ailego::Monotime::MilliSeconds() - start_time); LOG_INFO("DiskAnnBuilder::dump"); diff --git a/src/core/algorithm/diskann/diskann_builder.h b/src/core/algorithm/diskann/diskann_builder.h index 25e7310a0..dec19ec72 100644 --- a/src/core/algorithm/diskann/diskann_builder.h +++ b/src/core/algorithm/diskann/diskann_builder.h @@ -70,6 +70,9 @@ class DiskAnnBuilder : public IndexBuilder { int calculate_pq_chunk_num(); + //! Publish the first worker failure before waking the coordinator. + bool record_worker_error(int error_code); + double get_memory_in_bytes(double search_ram_budget) { return search_ram_budget * 1024 * 1024 * 1024; } diff --git a/src/core/algorithm/diskann/diskann_builder_entity.cc b/src/core/algorithm/diskann/diskann_builder_entity.cc index 393d01173..92cf191aa 100644 --- a/src/core/algorithm/diskann/diskann_builder_entity.cc +++ b/src/core/algorithm/diskann/diskann_builder_entity.cc @@ -13,21 +13,40 @@ // limitations under the License. #include "diskann_builder_entity.h" +#include +#include #include +#include +#include +#include +#include #include "diskann_algorithm.h" #include "diskann_util.h" namespace zvec { namespace core { +namespace { + +void update_atomic_max(std::atomic *value, uint32_t candidate) { + uint32_t current = value->load(std::memory_order_relaxed); + while (current < candidate && + !value->compare_exchange_weak(current, candidate, + std::memory_order_relaxed, + std::memory_order_relaxed)) { + } +} + +} // namespace + void DiskAnnBuilderEntity::clear() { max_degree_ = 0; list_size_ = 0; memory_limit_ = 0; num_threads_ = 0; max_build_degree_ = 0; - max_observed_degree_ = 0; - neighbor_size_ = 0; + max_observed_degree_.store(0, std::memory_order_relaxed); + neighbor_stride_ = 0; mem_index_file_.clear(); index_path_prefix_.clear(); vectors_buffer_.clear(); @@ -47,6 +66,17 @@ int DiskAnnBuilderEntity::init(const IndexMeta &meta, uint32_t max_degree, uint32_t list_size, double memory_limit, uint32_t build_threads) { clear(); + const double max_build_degree = + std::ceil(static_cast(max_degree) * + static_cast(kDefaultGraphSlackFactor)); + if (max_degree == 0 || list_size == 0 || + max_build_degree > + static_cast((std::numeric_limits::max)() - 1U)) { + LOG_ERROR("Invalid DiskAnn graph parameters: max_degree=%u list_size=%u", + max_degree, list_size); + return IndexError_InvalidArgument; + } + meta_ = meta; max_degree_ = max_degree; @@ -56,34 +86,56 @@ int DiskAnnBuilderEntity::init(const IndexMeta &meta, uint32_t max_degree, num_threads_ = build_threads; - max_build_degree_ = max_degree_ * kDefaultGraphSlackFactor; + max_build_degree_ = static_cast(max_build_degree); - neighbor_size_ = sizeof(uint32_t) + max_build_degree_ * sizeof(diskann_id_t); + // Store the neighbor count and ids in typed storage. Besides avoiding + // repeated byte conversions, this guarantees the alignment required by + // callers that consume the returned diskann_id_t pointer. + neighbor_stride_ = max_build_degree_ + 1; return 0; } -int DiskAnnBuilderEntity::reserve_space(uint32_t docs) { - vectors_buffer_.reserve(meta_.element_size() * docs); - keys_buffer_.reserve(sizeof(diskann_key_t) * docs); - neighbors_buffer_.reserve(neighbor_size_ * docs); +int DiskAnnBuilderEntity::reserve_space(size_t docs) { + if (docs == 0 || docs > static_cast(kInvalidId)) { + LOG_ERROR("Invalid DiskAnn document count: %zu", docs); + return IndexError_InvalidLength; + } + + const size_t element_size = meta_.element_size(); + if (element_size == 0 || + docs > (std::numeric_limits::max)() / element_size || + docs > (std::numeric_limits::max)() / sizeof(diskann_key_t) || + neighbor_stride_ == 0 || + docs > (std::numeric_limits::max)() / neighbor_stride_) { + LOG_ERROR("DiskAnn builder buffer size overflows: docs=%zu", docs); + return IndexError_InvalidLength; + } + + try { + vectors_buffer_.reserve(element_size * docs); + keys_buffer_.reserve(sizeof(diskann_key_t) * docs); + neighbors_buffer_.reserve(static_cast(neighbor_stride_) * docs); + } catch (const std::bad_alloc &) { + return IndexError_NoMemory; + } catch (const std::length_error &) { + return IndexError_InvalidLength; + } return 0; } int DiskAnnBuilderEntity::add_vector(diskann_key_t key, const void *vec) { + if (vec == nullptr) { + LOG_ERROR("Cannot add a null vector to DiskAnn"); + return IndexError_InvalidArgument; + } vectors_buffer_.append(reinterpret_cast(vec), meta_.element_size()); keys_buffer_.append(reinterpret_cast(&key), sizeof(key)); - uint32_t neighbor_cnt = 0; - - std::vector neighbor(max_build_degree_, 0); - - neighbors_buffer_.append(reinterpret_cast(&neighbor_cnt), - sizeof(uint32_t)); - neighbors_buffer_.append(reinterpret_cast(neighbor.data()), - sizeof(diskann_id_t) * max_build_degree_); + neighbors_buffer_.push_back(0); + neighbors_buffer_.resize(neighbors_buffer_.size() + max_build_degree_, 0); (*mutable_doc_cnt())++; @@ -97,9 +149,9 @@ const void *DiskAnnBuilderEntity::get_vector(diskann_id_t id) const { diskann_key_t DiskAnnBuilderEntity::get_key(diskann_id_t id) const { size_t offset = (size_t)id * sizeof(diskann_key_t); - - return *( - reinterpret_cast(keys_buffer_.data() + offset)); + diskann_key_t key = kInvalidKey; + memcpy(&key, keys_buffer_.data() + offset, sizeof(key)); + return key; } //! Get vector local id by key @@ -110,58 +162,46 @@ diskann_id_t DiskAnnBuilderEntity::get_id(diskann_key_t /*key*/) const { std::pair DiskAnnBuilderEntity::get_neighbors( diskann_id_t id) const { - size_t offset = (size_t)id * neighbor_size_; - - const uint8_t *start_ptr = - reinterpret_cast(neighbors_buffer_.data()) + offset; - - uint32_t neighbor_cnt = *(reinterpret_cast(start_ptr)); - - const diskann_id_t *neighbors = - reinterpret_cast(start_ptr + sizeof(uint32_t)); - - return std::make_pair(neighbor_cnt, neighbors); + const size_t offset = static_cast(id) * neighbor_stride_; + return std::make_pair(neighbors_buffer_[offset], + neighbors_buffer_.data() + offset + 1); } int DiskAnnBuilderEntity::set_neighbors( diskann_id_t id, const std::vector &neighbor_ids) { - size_t offset = (size_t)id * neighbor_size_; - - uint8_t *start_ptr = - reinterpret_cast(&neighbors_buffer_[0]) + offset; - - uint32_t neighbor_cnt = neighbor_ids.size(); - - memcpy(start_ptr + sizeof(uint32_t), neighbor_ids.data(), - sizeof(diskann_id_t) * neighbor_cnt); - memcpy(start_ptr, &neighbor_cnt, sizeof(uint32_t)); - - if (max_observed_degree_ < neighbor_cnt) { - max_observed_degree_ = neighbor_cnt; + if (id >= doc_cnt() || neighbor_ids.size() > max_build_degree_) { + LOG_ERROR("Invalid DiskAnn neighbor update: id=%u count=%zu", id, + neighbor_ids.size()); + return IndexError_OutOfRange; } + const size_t offset = static_cast(id) * neighbor_stride_; + const uint32_t neighbor_cnt = static_cast(neighbor_ids.size()); + std::copy(neighbor_ids.begin(), neighbor_ids.end(), + neighbors_buffer_.begin() + offset + 1); + neighbors_buffer_[offset] = neighbor_cnt; + + update_atomic_max(&max_observed_degree_, neighbor_cnt); return 0; } int DiskAnnBuilderEntity::add_neighbor(diskann_id_t id, diskann_id_t neighbor_id) { - size_t offset = (size_t)id * neighbor_size_; - - uint8_t *start_ptr = - reinterpret_cast(&neighbors_buffer_[0]) + offset; - - uint32_t neighbor_cnt = *reinterpret_cast(start_ptr); - - memcpy(start_ptr + sizeof(uint32_t) + sizeof(diskann_id_t) * neighbor_cnt, - &neighbor_id, sizeof(diskann_id_t)); - - neighbor_cnt += 1; - - memcpy(start_ptr, &neighbor_cnt, sizeof(uint32_t)); - - if (max_observed_degree_ < neighbor_cnt) { - max_observed_degree_ = neighbor_cnt; + if (id >= doc_cnt()) { + LOG_ERROR("Invalid DiskAnn node id: %u", id); + return IndexError_OutOfRange; + } + const size_t offset = static_cast(id) * neighbor_stride_; + uint32_t &neighbor_cnt = neighbors_buffer_[offset]; + if (neighbor_cnt >= max_build_degree_) { + LOG_ERROR("DiskAnn neighbor list is full: id=%u count=%u", id, + neighbor_cnt); + return IndexError_IndexFull; } + neighbors_buffer_[offset + 1 + neighbor_cnt] = neighbor_id; + ++neighbor_cnt; + + update_atomic_max(&max_observed_degree_, neighbor_cnt); return 0; } @@ -350,19 +390,50 @@ int DiskAnnBuilderEntity::dump_key_segment( return 0; } +int DiskAnnBuilderEntity::build_key_mapping( + std::vector *mapping) const { + mapping->resize(doc_cnt()); + auto get_key = [this](diskann_id_t id) { + diskann_key_t key = kInvalidKey; + memcpy( + &key, + keys_buffer_.data() + static_cast(id) * sizeof(diskann_key_t), + sizeof(key)); + return key; + }; + + std::iota(mapping->begin(), mapping->end(), 0U); + std::sort( + mapping->begin(), mapping->end(), + [&](diskann_id_t i, diskann_id_t j) { return get_key(i) < get_key(j); }); + + for (size_t i = 1; i < mapping->size(); ++i) { + const diskann_key_t previous_key = get_key((*mapping)[i - 1]); + const diskann_key_t current_key = get_key((*mapping)[i]); + if (current_key != kInvalidKey && current_key == previous_key) { + LOG_ERROR("Duplicate DiskAnn vector key: %llu", + static_cast(current_key)); + return IndexError_Exist; + } + } + return 0; +} + int DiskAnnBuilderEntity::dump_key_mapping_segment( const IndexDumper::Pointer &dumper) const { - std::vector mapping(doc_cnt()); - - const diskann_key_t *keys = reinterpret_cast( - const_cast(keys_buffer_.data())); - - std::iota(mapping.begin(), mapping.end(), 0U); - std::sort(mapping.begin(), mapping.end(), - [&](diskann_id_t i, diskann_id_t j) { return keys[i] < keys[j]; }); + std::vector mapping; + const int ret = build_key_mapping(&mapping); + if (ret != 0) { + return ret; + } + return dump_key_mapping_segment(dumper, mapping); +} - size_t size = mapping.size() * sizeof(diskann_id_t); - int64_t ret = +int DiskAnnBuilderEntity::dump_key_mapping_segment( + const IndexDumper::Pointer &dumper, + const std::vector &mapping) const { + const size_t size = mapping.size() * sizeof(diskann_id_t); + const int64_t ret = dump_segment(dumper, kDiskAnnKeyMappingSegmentId, mapping.data(), size); if (ret != 0) { @@ -400,11 +471,19 @@ int DiskAnnBuilderEntity::dump_entrypoint_segment( return 0; } -int DiskAnnBuilderEntity::dump(IndexHolder::Pointer holder, IndexMeta &meta, +int DiskAnnBuilderEntity::dump(IndexMeta &meta, const IndexDumper::Pointer &dumper) { - uint64_t doc_cnt = holder->count(); + const uint64_t doc_cnt = this->doc_cnt(); + const uint32_t max_observed_degree = + max_observed_degree_.load(std::memory_order_acquire); + std::vector key_mapping; + int ret = build_key_mapping(&key_mapping); + if (ret != 0) { + LOG_ERROR("Failed to build key mapping"); + return ret; + } uint64_t max_node_size = - (uint64_t)max_observed_degree_ * sizeof(diskann_id_t) + sizeof(uint32_t) + + (uint64_t)max_observed_degree * sizeof(diskann_id_t) + sizeof(uint32_t) + meta_.element_size(); uint64_t node_per_sector = DiskAnnUtil::kSectorSize / @@ -413,18 +492,14 @@ int DiskAnnBuilderEntity::dump(IndexHolder::Pointer holder, IndexMeta &meta, std::string node_buf; node_buf.resize(max_node_size); - diskann_id_t *neighbor_buf = - (diskann_id_t *)(node_buf.data() + (meta_.element_size()) + - sizeof(uint32_t)); - LOG_INFO( "Dump Data, medoid: %zu, max node size: %zu, node per sector: %zu, " "max observed degree: %zu", (size_t)medoid(), (size_t)max_node_size, (size_t)node_per_sector, - (size_t)max_observed_degree_); + (size_t)max_observed_degree); // write a dummy segment to make data align - int ret = dump_dummy_segment(dumper); + ret = dump_dummy_segment(dumper); if (ret != 0) { LOG_ERROR("Dump dummy segment failed"); @@ -436,13 +511,6 @@ int DiskAnnBuilderEntity::dump(IndexHolder::Pointer holder, IndexMeta &meta, uint32_t crc = 0U; size_t len = 0; - // no need to write first sector - auto iter = holder->create_iterator(); - if (!iter) { - LOG_ERROR("Create iterator for holder failed"); - return IndexError_Runtime; - } - uint64_t index_size = 0; uint32_t neighbor_num; if (node_per_sector > 0) { @@ -469,27 +537,19 @@ int DiskAnnBuilderEntity::dump(IndexHolder::Pointer holder, IndexMeta &meta, auto neighbors = get_neighbors(cur_node_id); neighbor_num = neighbors.first; - ailego_assert(neighbor_num > 0); - ailego_assert(neighbor_num <= max_observed_degree_); - - memcpy(&(neighbor_buf[0]), neighbors.second, - neighbors.first * sizeof(diskann_id_t)); - - if (iter->is_valid()) { - const void *vec = iter->data(); - memcpy(&(node_buf[0]), vec, meta.element_size()); + ailego_assert(neighbor_num > 0 || doc_cnt == 1); + ailego_assert(neighbor_num <= max_observed_degree); - iter->next(); - } else { - return IndexError_Runtime; - } + const void *vec = get_vector(cur_node_id); + memcpy(&(node_buf[0]), vec, meta.element_size()); // write neighbor num - *(uint32_t *)(node_buf.data() + meta_.element_size()) = neighbor_num; + memcpy(node_buf.data() + meta_.element_size(), &neighbor_num, + sizeof(neighbor_num)); // write neighbor buffer memcpy(&(node_buf[0]) + meta_.element_size() + sizeof(uint32_t), - neighbor_buf, neighbor_num * sizeof(diskann_id_t)); + neighbors.second, neighbor_num * sizeof(diskann_id_t)); // get offset into sector_buf char *sector_node_buf = §or_buf[sector_node_id * max_node_size]; @@ -537,29 +597,19 @@ int DiskAnnBuilderEntity::dump(IndexHolder::Pointer holder, IndexMeta &meta, auto neighbors = get_neighbors(i); neighbor_num = neighbors.first; - ailego_assert(neighbor_num > 0); - ailego_assert(neighbor_num <= max_observed_degree_); + ailego_assert(neighbor_num > 0 || doc_cnt == 1); + ailego_assert(neighbor_num <= max_observed_degree); - // read node's nhood - memcpy((char *)neighbor_buf, neighbors.second, - neighbor_num * sizeof(diskann_id_t)); - - if (iter->is_valid()) { - const void *vec = iter->data(); - memcpy(&(multisector_buf[0]), vec, meta.element_size()); - - iter->next(); - } else { - return IndexError_Runtime; - } + const void *vec = get_vector(static_cast(i)); + memcpy(&(multisector_buf[0]), vec, meta.element_size()); // write neighbor - *(uint32_t *)(&(multisector_buf[0]) + meta_.element_size()) = - neighbor_num; + memcpy(&(multisector_buf[0]) + meta_.element_size(), &neighbor_num, + sizeof(neighbor_num)); // write nhood next memcpy(&(multisector_buf[0]) + meta_.element_size() + sizeof(uint32_t), - neighbor_buf, neighbor_num * sizeof(diskann_id_t)); + neighbors.second, neighbor_num * sizeof(diskann_id_t)); // flush sector to disk len = dumper->write(multisector_buf.data(), @@ -604,7 +654,7 @@ int DiskAnnBuilderEntity::dump(IndexHolder::Pointer holder, IndexMeta &meta, meta_header_.ndims = meta_.dimension(); meta_header_.medoid = medoid(); meta_header_.max_node_size = max_node_size; - meta_header_.max_degree = max_observed_degree_; + meta_header_.max_degree = max_observed_degree; meta_header_.node_per_sector = node_per_sector; meta_header_.vamana_frozen_num = 0; meta_header_.vamana_frozen_loc = medoid(); @@ -644,7 +694,7 @@ int DiskAnnBuilderEntity::dump(IndexHolder::Pointer holder, IndexMeta &meta, } // dump key mapping - ret = dump_key_mapping_segment(dumper); + ret = dump_key_mapping_segment(dumper, key_mapping); if (ret != 0) { LOG_ERROR("Dump key mapping segment failed"); diff --git a/src/core/algorithm/diskann/diskann_builder_entity.h b/src/core/algorithm/diskann/diskann_builder_entity.h index 653ea377e..e777cd15e 100644 --- a/src/core/algorithm/diskann/diskann_builder_entity.h +++ b/src/core/algorithm/diskann/diskann_builder_entity.h @@ -13,6 +13,7 @@ // limitations under the License. #pragma once +#include #include #include #include "diskann_entity.h" @@ -49,8 +50,7 @@ class DiskAnnBuilderEntity : public DiskAnnEntity { int init(const IndexMeta &meta, uint32_t max_degree, uint32_t list_size, double memory_limit, uint32_t build_threads); - int dump(IndexHolder::Pointer holder, IndexMeta &meta, - const IndexDumper::Pointer &dumper); + int dump(IndexMeta &meta, const IndexDumper::Pointer &dumper); int64_t dump_segment(const IndexDumper::Pointer &dumper, const std::string &segment_id, const void *data, @@ -62,7 +62,7 @@ class DiskAnnBuilderEntity : public DiskAnnEntity { int dump_entrypoint_segment(const IndexDumper::Pointer &dumper) const; int dump_key_segment(const IndexDumper::Pointer &dumper) const; - int reserve_space(uint32_t docs); + int reserve_space(size_t docs); std::vector &pq_full_pivot_data() { return pq_full_pivot_data_; @@ -86,15 +86,15 @@ class DiskAnnBuilderEntity : public DiskAnnEntity { double memory_limit_{0}; uint32_t num_threads_{0}; uint32_t max_build_degree_{0}; - uint32_t max_observed_degree_{0}; - uint32_t neighbor_size_{0}; + std::atomic max_observed_degree_{0}; + uint32_t neighbor_stride_{0}; std::string mem_index_file_{""}; std::string index_path_prefix_{""}; std::string vectors_buffer_{}; std::string keys_buffer_{}; - std::string neighbors_buffer_{}; + std::vector neighbors_buffer_{}; std::vector entrypoints_{}; IndexMeta meta_; @@ -103,6 +103,10 @@ class DiskAnnBuilderEntity : public DiskAnnEntity { std::vector pq_centroid_; std::vector pq_chunk_offsets_; std::vector block_compressed_data_; + + int build_key_mapping(std::vector *mapping) const; + int dump_key_mapping_segment(const IndexDumper::Pointer &dumper, + const std::vector &mapping) const; }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_context.cc b/src/core/algorithm/diskann/diskann_context.cc index 220979420..38d11fd57 100644 --- a/src/core/algorithm/diskann/diskann_context.cc +++ b/src/core/algorithm/diskann/diskann_context.cc @@ -14,6 +14,7 @@ #include "diskann_context.h" #include +#include #include "diskann_params.h" #include "diskann_pq_table.h" #include "diskann_util.h" @@ -28,22 +29,78 @@ DiskAnnContext::DiskAnnContext(const IndexMeta &meta, dc_(entity.get(), measure, meta.dimension()), entity_{entity} {} +DiskAnnContext::Pointer DiskAnnContext::create_fetch_context( + const IndexMeta &meta, const IndexMetric::Pointer &measure, + const DiskAnnEntity::Pointer &entity) { + if (!measure || !entity) { + return nullptr; + } + + Pointer context(new (std::nothrow) DiskAnnContext(meta, measure, entity)); + if (!context || + context->init(kFetchContext, entity->max_degree(), entity->pq_chunk_num(), + meta.element_size()) != 0) { + return nullptr; + } + return context; +} + +int DiskAnnContext::resize_fetch_sector_buffer( + const DiskAnnEntity::Pointer &entity) { + if (!entity) { + LOG_ERROR("Cannot size a DiskAnn fetch buffer without an entity"); + return IndexError_InvalidArgument; + } + + const uint64_t sector_num_per_node = + entity->node_per_sector() > 0 + ? 1 + : DiskAnnUtil::div_round_up(entity->max_node_size(), + DiskAnnUtil::kSectorSize); + if (sector_num_per_node == 0 || + sector_num_per_node > DiskAnnUtil::kMaxSectorReadNum) { + LOG_ERROR("Invalid DiskAnn fetch sector count: %lu", + static_cast(sector_num_per_node)); + return IndexError_InvalidArgument; + } + + const size_t required_size = + static_cast(sector_num_per_node) * DiskAnnUtil::kSectorSize; + if (sector_buffer_ != nullptr && sector_buffer_size_ == required_size) { + return 0; + } + + void *replacement = nullptr; + DiskAnnUtil::alloc_aligned(&replacement, required_size, + DiskAnnUtil::kSectorSize); + if (!replacement) { + LOG_ERROR("Failed to allocate DiskAnn fetch buffer"); + return IndexError_NoMemory; + } + + DiskAnnUtil::free_aligned(sector_buffer_); + sector_buffer_ = replacement; + sector_buffer_size_ = required_size; + return 0; +} + int DiskAnnContext::init(ContextType type, uint32_t graph_degree, uint32_t pq_chunk_num, uint32_t element_size) { if (!entity_ || element_size == 0) { LOG_ERROR("Invalid DiskAnn context parameters"); return IndexError_InvalidArgument; } - type_ = type; element_size_ = element_size; pq_chunk_num_ = pq_chunk_num; - DiskAnnUtil::alloc_aligned((void **)&query_, element_size_, 32); - DiskAnnUtil::alloc_aligned((void **)&query_rotated_, element_size_, 32); - if (!query_ || !query_rotated_) { - LOG_ERROR("Failed to allocate DiskAnn query buffers"); - return IndexError_NoMemory; + if (type != kFetchContext) { + DiskAnnUtil::alloc_aligned((void **)&query_, element_size_, 32); + DiskAnnUtil::alloc_aligned((void **)&query_rotated_, element_size_, 32); + if (!query_ || !query_rotated_) { + LOG_ERROR("Failed to allocate DiskAnn query buffers"); + return IndexError_NoMemory; + } } int ret; @@ -58,7 +115,8 @@ int DiskAnnContext::init(ContextType type, uint32_t graph_degree, break; case kSearcherContext: - if (graph_degree == 0 || pq_chunk_num_ == 0) { + if (pq_chunk_num_ == 0 || + (graph_degree == 0 && entity_->doc_cnt() != 1)) { LOG_ERROR("Invalid DiskAnn search context dimensions"); return IndexError_InvalidArgument; } @@ -76,13 +134,14 @@ int DiskAnnContext::init(ContextType type, uint32_t graph_degree, 256); DiskAnnUtil::alloc_aligned( (void **)&pq_coord_buffer_, - static_cast(graph_degree) * pq_chunk_num_ * sizeof(uint8_t), + static_cast(std::max(graph_degree, 1U)) * pq_chunk_num_ * + sizeof(uint8_t), 256); DiskAnnUtil::alloc_aligned((void **)&coord_buffer_, element_size_, 256); - DiskAnnUtil::alloc_aligned( - (void **)§or_buffer_, - DiskAnnUtil::kMaxSectorReadNum * DiskAnnUtil::kSectorSize, - DiskAnnUtil::kSectorSize); + sector_buffer_size_ = static_cast(DiskAnnUtil::kMaxSectorReadNum * + DiskAnnUtil::kSectorSize); + DiskAnnUtil::alloc_aligned((void **)§or_buffer_, sector_buffer_size_, + DiskAnnUtil::kSectorSize); if (!pq_table_dist_buffer_ || !pq_coord_buffer_ || !coord_buffer_ || !sector_buffer_) { LOG_ERROR("Failed to allocate DiskAnn search buffers"); @@ -96,6 +155,19 @@ int DiskAnnContext::init(ContextType type, uint32_t graph_degree, } break; + case kFetchContext: + ret = resize_fetch_sector_buffer(entity_); + if (ret != 0) { + return ret; + } + + ret = setup_io_ctx(io_ctx_); + if (ret != 0) { + LOG_ERROR("setup fetch io ctx error, ret=%d", ret); + return ret; + } + break; + default: LOG_ERROR("Init context failed"); return IndexError_Runtime; @@ -105,22 +177,30 @@ int DiskAnnContext::init(ContextType type, uint32_t graph_degree, } DiskAnnContext::~DiskAnnContext() { - free(query_); - free(query_rotated_); - free(pq_table_dist_buffer_); - free(pq_coord_buffer_); - free(coord_buffer_); - free(sector_buffer_); - - if (type_ == kSearcherContext) { + // The sector buffer may still be the destination of an overlapped read if a + // query exits early. Cancel and wait for every request before releasing any + // memory that the I/O context can reference. + if (type_ == kSearcherContext || type_ == kFetchContext) { destroy_io_ctx(io_ctx_); } + + visit_filter_.destroy(); + DiskAnnUtil::free_aligned(query_); + DiskAnnUtil::free_aligned(query_rotated_); + DiskAnnUtil::free_aligned(pq_table_dist_buffer_); + DiskAnnUtil::free_aligned(pq_coord_buffer_); + DiskAnnUtil::free_aligned(coord_buffer_); + DiskAnnUtil::free_aligned(sector_buffer_); } int DiskAnnContext::update(const ailego::Params ¶ms) { uint32_t list_size = list_size_; params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size); - list_size_ = list_size; + if (list_size == 0) { + LOG_ERROR("list_size must be positive"); + return IndexError_InvalidArgument; + } + set_list_size(list_size); return 0; } @@ -128,7 +208,7 @@ int DiskAnnContext::update_context(ContextType type, const IndexMeta &meta, const IndexMetric::Pointer &measure, const DiskAnnEntity::Pointer &entity, uint32_t magic_num) { - if (ailego_unlikely(type != type_)) { + if (ailego_unlikely(type != static_cast(type_))) { LOG_ERROR( "DiskAnnContext does not support shared by different type, " "src=%u dst=%u", @@ -146,6 +226,14 @@ int DiskAnnContext::update_context(ContextType type, const IndexMeta &meta, case kSearcherContext: break; + case kFetchContext: { + const int ret = resize_fetch_sector_buffer(entity); + if (ret != 0) { + return ret; + } + break; + } + case kReducerContext: break; @@ -155,6 +243,7 @@ int DiskAnnContext::update_context(ContextType type, const IndexMeta &meta, } entity_ = entity; + set_list_size(requested_list_size_); update_index_metric(measure); dc_.update(entity_.get(), measure, meta.dimension()); magic_ = magic_num; diff --git a/src/core/algorithm/diskann/diskann_context.h b/src/core/algorithm/diskann/diskann_context.h index 5934cc03c..0ca4557b7 100644 --- a/src/core/algorithm/diskann/diskann_context.h +++ b/src/core/algorithm/diskann/diskann_context.h @@ -13,6 +13,7 @@ // limitations under the License. #pragma once +#include #include #include "utility/topk_result_builder.h" #include "diskann_dist_calculator.h" @@ -47,7 +48,8 @@ class DiskAnnContext : public IndexContext, kUnknownContext = 0, kSearcherContext = 1, kBuilderContext = 2, - kReducerContext = 3 + kReducerContext = 3, + kFetchContext = 4 }; //! Construct @@ -57,6 +59,11 @@ class DiskAnnContext : public IndexContext, //! Destructor virtual ~DiskAnnContext(); + //! Create a lightweight context for reading vectors by id. + static Pointer create_fetch_context(const IndexMeta &meta, + const IndexMetric::Pointer &measure, + const DiskAnnEntity::Pointer &entity); + public: //! Init int init(ContextType type, uint32_t graph_degree, uint32_t pq_chunk_num, @@ -139,7 +146,13 @@ class DiskAnnContext : public IndexContext, } void set_list_size(uint32_t list_size) { - list_size_ = list_size; + requested_list_size_ = list_size; + if (entity_ && entity_->doc_cnt() > 0) { + list_size_ = static_cast( + std::min(list_size, entity_->doc_cnt())); + } else { + list_size_ = list_size; + } } void set_fetch_vector(bool v) override { @@ -155,6 +168,10 @@ class DiskAnnContext : public IndexContext, return list_size_; } + inline uint32_t requested_list_size() const { + return requested_list_size_; + } + inline void reset_query(const void *query) { memcpy(query_, query, element_size_); memcpy(query_rotated_, query, element_size_); @@ -190,10 +207,18 @@ class DiskAnnContext : public IndexContext, return sector_buffer_; } + inline size_t sector_buffer_size() const { + return sector_buffer_size_; + } + inline IOContext &io_ctx() { return io_ctx_; } + ContextType context_type() const { + return static_cast(type_); + } + inline void resize_results(size_t size) { if (group_by_search()) { group_results_.resize(size); @@ -236,7 +261,7 @@ class DiskAnnContext : public IndexContext, group_num_ = rhs.group_num_; group_topk_ = rhs.group_topk_; group_topk_heaps_.clear(); - list_size_ = rhs.list_size_; + set_list_size(rhs.requested_list_size_); fetch_vector_ = rhs.fetch_vector_; debug_mode_ = rhs.debug_mode_; } @@ -293,6 +318,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_; + set_list_size(other.requested_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; @@ -339,13 +379,19 @@ class DiskAnnContext : public IndexContext, void emplace_result_doc(IndexDocumentList &docs, diskann_id_t id, float score, const VectorInfo &info) { + const diskann_key_t key = entity_->get_key(id); + if (key == kInvalidKey) { + return; + } if (fetch_vector_) { - docs.emplace_back(entity_->get_key(id), score, id, info.vec_); + docs.emplace_back(key, score, id, info.vec_); } else { - docs.emplace_back(entity_->get_key(id), score, id); + docs.emplace_back(key, score, id); } } + int resize_fetch_sector_buffer(const DiskAnnEntity::Pointer &entity); + private: constexpr static uint32_t kInvalidMgic = -1U; @@ -361,6 +407,7 @@ class DiskAnnContext : public IndexContext, uint32_t element_size_{0}; uint32_t element_rotated_size_{0}; uint32_t list_size_{0}; + uint32_t requested_list_size_{0}; TopkHeap topk_heap_{}; @@ -368,7 +415,7 @@ class DiskAnnContext : public IndexContext, uint32_t group_num_{0}; std::map group_topk_heaps_{}; - IOContext io_ctx_{0}; + IOContext io_ctx_{}; SearchStats query_stats_; float *pq_table_dist_buffer_{nullptr}; @@ -377,6 +424,7 @@ class DiskAnnContext : public IndexContext, void *query_rotated_{nullptr}; void *coord_buffer_{nullptr}; void *sector_buffer_{nullptr}; + size_t sector_buffer_size_{0}; std::vector results_{}; std::vector group_results_{}; diff --git a/src/core/algorithm/diskann/diskann_entity.h b/src/core/algorithm/diskann/diskann_entity.h index bd5ce5142..7babd67e6 100644 --- a/src/core/algorithm/diskann/diskann_entity.h +++ b/src/core/algorithm/diskann/diskann_entity.h @@ -24,7 +24,7 @@ using diskann_key_t = uint64_t; using diskann_id_t = uint32_t; constexpr diskann_id_t kInvalidId = static_cast(-1); -constexpr diskann_key_t kInvalidKey = static_cast(-1); +constexpr diskann_key_t kInvalidKey = static_cast(-1); struct VectorInfo { float dist_; @@ -178,6 +178,14 @@ class DiskAnnEntity { return meta_header_.doc_cnt; } + uint64_t dimension() const { + return meta_header_.ndims; + } + + uint64_t index_size() const { + return meta_header_.index_size; + } + uint64_t *mutable_doc_cnt() { return &meta_header_.doc_cnt; } diff --git a/src/core/algorithm/diskann/diskann_file_reader.cc b/src/core/algorithm/diskann/diskann_file_reader.cc index 681901e34..a04c8e8d8 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.cc +++ b/src/core/algorithm/diskann/diskann_file_reader.cc @@ -14,19 +14,26 @@ #include "diskann_file_reader.h" #include -#include +#include #include #include #include +#include #include #include +#include #include #include #include +#include #include #include #include +#if defined(_WIN32) || defined(_WIN64) +#include +#endif #if defined(__APPLE__) || defined(__MACH__) +#include #include #include #endif @@ -42,7 +49,7 @@ static std::once_flag g_io_backend_log_once; static void log_diskann_io_backend(ailego::IOBackendType type) { #if (defined(__linux) || defined(__linux__) || defined(__APPLE__) || \ - defined(__MACH__)) + defined(__MACH__) || defined(_WIN32) || defined(_WIN64)) std::call_once(g_io_backend_log_once, [type]() { #if (defined(__linux) || defined(__linux__)) if (type == ailego::IOBackendType::kPread) { @@ -56,6 +63,9 @@ static void log_diskann_io_backend(ailego::IOBackendType type) { LOG_INFO("DiskAnn: I/O backend '%s' loaded — async I/O enabled.", ailego::IOBackendTypeName(type)); } +#elif defined(_WIN32) || defined(_WIN64) + LOG_INFO("DiskAnn: I/O backend '%s' — asynchronous I/O enabled.", + ailego::IOBackendTypeName(type)); #else LOG_INFO("DiskAnn: I/O backend '%s' — synchronous I/O enabled.", ailego::IOBackendTypeName(type)); @@ -79,6 +89,148 @@ void log_diskann_io_backend() { log_diskann_io_backend(ailego::IOBackend::Instance().available()); } +#if defined(_WIN32) || defined(_WIN64) +// The reader's stable handle owns the selected file object across lazy I/O +// batches. Allow its path to be deleted or atomically replaced, but keep +// in-place writes blocked while reads are active. +static constexpr DWORD kDiskAnnFileShareMode = + FILE_SHARE_READ | FILE_SHARE_DELETE; +// Every handle kept beside an IOCP handle must bypass the system cache. A +// buffered handle to the same file object substantially degrades 4 KiB random +// reads even when the actual ReadFile calls use a separate unbuffered handle. +static constexpr DWORD kDiskAnnStableHandleFlags = FILE_FLAG_NO_BUFFERING; + +// Threads stay associated with an IOCP after dequeuing a completion. An +// IOContext can move between serialized callers, so a concurrency limit of one +// can permanently starve a later caller while the previous thread remains +// runnable. Callers still serialize each context, and outstanding_count +// rejects overlapping batches; keep the private port effectively unthrottled. +static constexpr DWORD kDiskAnnIoCompletionConcurrency = + static_cast(MAXLONG); + +// An IOContext may be reused with different reader instances. Assign every +// successfully opened file object a process-wide identity so two readers for +// the same path cannot accidentally share an IOCP handle to stale contents. +static std::atomic g_next_windows_file_identity{1}; + +static uint64_t next_windows_file_identity() { + uint64_t identity = + g_next_windows_file_identity.load(std::memory_order_relaxed); + while (identity != 0) { + const uint64_t next = + identity == std::numeric_limits::max() ? 0 : identity + 1; + if (g_next_windows_file_identity.compare_exchange_weak( + identity, next, std::memory_order_relaxed, + std::memory_order_relaxed)) { + return identity; + } + } + return 0; +} + +// Cancel and reap every request that may still reference caller-owned buffers. +// Closing the file or completion port before the cancellation packets have +// arrived would allow the kernel to keep writing into buffers already returned +// to the caller. +static void close_windows_io_handles(IOContext ctx) { + if (ctx == nullptr) { + return; + } + + size_t active_remaining = static_cast( + std::count_if(ctx->active_requests.begin(), ctx->active_requests.end(), + [](uint8_t active) { return active != 0; })); + // The branches below are internal-state failures, not ordinary I/O errors. + // Returning after one of them could free an OVERLAPPED or destination + // buffer that the kernel still owns, so fail closed instead of risking UAF. + if (active_remaining != ctx->outstanding_count) { + LOG_FATAL("DiskAnn Windows I/O context lost track of outstanding requests"); + std::abort(); + } + if (active_remaining != 0) { + if (ctx->file_handle == INVALID_HANDLE_VALUE || + ctx->completion_port == nullptr) { + LOG_FATAL( + "Cannot drain DiskAnn overlapped requests without valid handles"); + std::abort(); + } + if (!::CancelIoEx(ctx->file_handle, nullptr)) { + DWORD error = ::GetLastError(); + if (error != ERROR_NOT_FOUND) { + LOG_WARN("CancelIoEx failed while draining DiskAnn I/O (error=%lu)", + error); + } + } + + // A file associated with an IOCP does not publish the final OVERLAPPED + // status until its completion packet is dequeued. Polling + // GetOverlappedResult() here can therefore return ERROR_IO_INCOMPLETE + // forever. The completion port is private to this context, so drain it + // until every active slot has produced its terminal packet. + while (active_remaining != 0) { + OVERLAPPED_ENTRY entries[MAX_IO_DEPTH]{}; + ULONG removed = 0; + const ULONG max_entries = + static_cast(std::min(active_remaining, MAX_IO_DEPTH)); + if (!::GetQueuedCompletionStatusEx(ctx->completion_port, entries, + max_entries, &removed, INFINITE, + FALSE) || + removed == 0) { + LOG_FATAL( + "GetQueuedCompletionStatusEx failed while draining DiskAnn I/O " + "(error=%lu)", + ::GetLastError()); + std::abort(); + } + + for (ULONG i = 0; i < removed; ++i) { + if (entries[i].lpCompletionKey != reinterpret_cast(ctx)) { + LOG_FATAL( + "DiskAnn teardown received a completion for another context"); + std::abort(); + } + + const uintptr_t address = + reinterpret_cast(entries[i].lpOverlapped); + const uintptr_t begin = reinterpret_cast(ctx->reqs.data()); + const uintptr_t span = ctx->reqs.size() * sizeof(OVERLAPPED); + if (address < begin || address >= begin + span || + (address - begin) % sizeof(OVERLAPPED) != 0) { + LOG_FATAL("DiskAnn teardown received an unknown OVERLAPPED request"); + std::abort(); + } + + const size_t index = (address - begin) / sizeof(OVERLAPPED); + if (ctx->active_requests[index] == 0) { + LOG_FATAL("DiskAnn teardown received a duplicate completion"); + std::abort(); + } + ctx->active_requests[index] = 0; + --active_remaining; + if (ctx->outstanding_count != 0) { + --ctx->outstanding_count; + } + } + } + ctx->outstanding_count = 0; + } + + if (ctx->file_handle != INVALID_HANDLE_VALUE) { + ::CloseHandle(ctx->file_handle); + ctx->file_handle = INVALID_HANDLE_VALUE; + } + + if (ctx->completion_port != nullptr) { + ::CloseHandle(ctx->completion_port); + ctx->completion_port = nullptr; + } + ctx->file_path.clear(); + ctx->file_identity = 0; + ctx->active_requests.fill(0); + ctx->outstanding_count = 0; +} +#endif + int setup_io_ctx(IOContext &ctx) { auto selected = ailego::IOBackend::Instance().available(); ctx = new (std::nothrow) IoBackend(); @@ -88,7 +240,10 @@ int setup_io_ctx(IOContext &ctx) { } ctx->type = selected; -#if (defined(__linux) || defined(__linux__)) +#if defined(_WIN32) || defined(_WIN64) + log_diskann_io_backend(ctx->type); + return 0; +#elif defined(__linux) || defined(__linux__) if (selected == ailego::IOBackendType::kPread) { log_diskann_io_backend(ctx->type); return 0; @@ -108,6 +263,9 @@ int setup_io_ctx(IOContext &ctx) { int ret = LibAioLoader::Instance().io_setup(MAX_EVENTS, &ctx->aio_ctx); if (ret == 0) { ctx->type = ailego::IOBackendType::kLibAio; + if (selected == ailego::IOBackendType::kIoUring) { + ailego::IOBackend::Instance().downgrade(ctx->type); + } log_diskann_io_backend(ctx->type); return 0; } @@ -117,6 +275,7 @@ int setup_io_ctx(IOContext &ctx) { // Priority 3: synchronous pread (always available). ctx->type = ailego::IOBackendType::kPread; + ailego::IOBackend::Instance().downgrade(ctx->type); #endif log_diskann_io_backend(ctx->type); return 0; @@ -127,7 +286,9 @@ int destroy_io_ctx(IOContext &ctx) { return 0; } -#if (defined(__linux) || defined(__linux__)) +#if defined(_WIN32) || defined(_WIN64) + close_windows_io_handles(ctx); +#elif defined(__linux) || defined(__linux__) if (ctx->type == ailego::IOBackendType::kIoUring) { ctx->ring.teardown(); } else if (ctx->type == ailego::IOBackendType::kLibAio && @@ -142,6 +303,10 @@ int destroy_io_ctx(IOContext &ctx) { return 0; } +#if !defined(_WIN32) && !defined(_WIN64) +static_assert(sizeof(off_t) >= sizeof(uint64_t), + "DiskAnn requires 64-bit POSIX file offsets"); + static int execute_one_pread(int fd, const AlignedRead &req) { auto *buf = static_cast(req.buf); uint64_t offset = req.offset; @@ -339,8 +504,8 @@ int execute_io_libaio(io_context_t &ctx, int fd, int execute_io(IOContext ctx, int fd, std::vector &read_reqs, uint64_t n_retries = 0) { #if (defined(__linux) || defined(__linux__)) - // Guard against null or sentinel contexts. - if (ctx == nullptr || ctx == (IOContext)-1) { + // A missing asynchronous context falls back to synchronous pread. + if (ctx == nullptr) { return execute_io_pread(fd, read_reqs); } // Dispatch based on the active backend. @@ -590,84 +755,126 @@ LinuxAlignedFileReader::LinuxAlignedFileReader() { } LinuxAlignedFileReader::~LinuxAlignedFileReader() { - deregister_all_threads(); if (file_desc >= 0) { ::close(file_desc); file_desc = -1; } } -IOContext &LinuxAlignedFileReader::get_ctx() { - std::unique_lock lk(ctx_mut); - auto it = ctx_map.find(std::this_thread::get_id()); - if (it == ctx_map.end()) { - LOG_ERROR("bad thread access; returning invalid IOContext"); - return this->bad_ctx; - } else { - return it->second; +static int duplicate_file_descriptor(int source_fd) { +#if defined(F_DUPFD_CLOEXEC) + return ::fcntl(source_fd, F_DUPFD_CLOEXEC, 0); +#else + int duplicate_fd = ::dup(source_fd); + if (duplicate_fd >= 0 && ::fcntl(duplicate_fd, F_SETFD, FD_CLOEXEC) == -1) { + const int saved_errno = errno; + ::close(duplicate_fd); + errno = saved_errno; + return -1; } + return duplicate_fd; +#endif } -void LinuxAlignedFileReader::register_thread() { - auto thread_id = std::this_thread::get_id(); - std::unique_lock lk(ctx_mut); - if (ctx_map.find(thread_id) != ctx_map.end()) { - LOG_ERROR("multiple calls to register_thread from the same thread"); - return; +#if defined(__linux__) || defined(__linux) +static int reopen_file_descriptor_with_direct_io(int source_fd) { + // dup()/F_DUPFD_CLOEXEC shares one open-file description with source_fd, so + // changing O_DIRECT through F_SETFL would also change the caller's buffered + // FileReadStorage handle. Reopening the procfs descriptor gives DiskAnn an + // independent open-file description while still referring to the exact + // inode captured during metadata loading (including an unlinked/replaced + // file). Some restricted environments do not mount procfs; callers fall + // back to a buffered duplicate in that case. + char fd_path[64]; + const int path_length = + std::snprintf(fd_path, sizeof(fd_path), "/proc/self/fd/%d", source_fd); + if (path_length <= 0 || static_cast(path_length) >= sizeof(fd_path)) { + errno = EINVAL; + return -1; } - IOContext ctx = nullptr; - int ret = setup_io_ctx(ctx); - if (ret != 0) { - LOG_ERROR("setup_io_ctx failed; returned: %d", ret); - lk.unlock(); - return; - } - if (ctx != nullptr) { - LOG_INFO("allocating ctx: %p", static_cast(ctx)); - } - ctx_map[thread_id] = ctx; - lk.unlock(); + int flags = O_RDONLY | O_DIRECT | O_LARGEFILE; +#if defined(O_CLOEXEC) + flags |= O_CLOEXEC; +#endif + return ::open(fd_path, flags); } +#endif -void LinuxAlignedFileReader::deregister_thread() { - auto thread_id = std::this_thread::get_id(); - IOContext ctx; +#if defined(__APPLE__) || defined(__MACH__) +static int reopen_macos_file_descriptor(const std::string &fname, + int source_fd) { + int flags = O_RDONLY; +#if defined(O_CLOEXEC) + flags |= O_CLOEXEC; +#endif + int reopened_fd = ::open(fname.c_str(), flags); + if (reopened_fd < 0) { + return -1; + } - { - std::lock_guard lk(ctx_mut); - auto it = ctx_map.find(thread_id); - if (it == ctx_map.end()) { - LOG_ERROR("deregister_thread: thread not registered"); - return; - } - ctx = it->second; - ctx_map.erase(it); +#if !defined(O_CLOEXEC) + if (::fcntl(reopened_fd, F_SETFD, FD_CLOEXEC) == -1) { + const int saved_errno = errno; + ::close(reopened_fd); + errno = saved_errno; + return -1; } +#endif - // Keep teardown outside the lock; async backends may block in syscalls. - destroy_io_ctx(ctx); - LOG_INFO("returned ctx from thread"); + struct stat source_stat {}; + struct stat reopened_stat {}; + if (::fstat(source_fd, &source_stat) == -1 || + ::fstat(reopened_fd, &reopened_stat) == -1) { + const int saved_errno = errno; + ::close(reopened_fd); + errno = saved_errno; + return -1; + } + if (source_stat.st_dev != reopened_stat.st_dev || + source_stat.st_ino != reopened_stat.st_ino) { + ::close(reopened_fd); + errno = ESTALE; + return -1; + } + return reopened_fd; } -void LinuxAlignedFileReader::deregister_all_threads() { - std::unique_lock lk(ctx_mut); - for (auto x = ctx_map.begin(); x != ctx_map.end(); x++) { - destroy_io_ctx(x->second); +static void configure_macos_reader(int file_desc, const std::string &fname) { + // macOS has no O_DIRECT. F_NOCACHE is its closest per-file equivalent: it + // asks the kernel to minimize caching for I/O through this descriptor. This + // is advisory rather than a guarantee that every read reaches the device. + // Disable read-ahead as well because DiskAnn performs random reads. + // + // Do not mmap the entire index and call msync(MS_INVALIDATE) here. That does + // not provide a reliable global cache eviction guarantee and makes open time + // and virtual-address usage scale with the size of the index. + if (::fcntl(file_desc, F_NOCACHE, 1) == -1) { + LOG_WARN( + "fcntl(F_NOCACHE) failed for %s (errno=%d: %s); reads will use " + "the page cache", + fname.c_str(), errno, ::strerror(errno)); + } else { + LOG_INFO("DiskAnn macOS: F_NOCACHE enabled for %s", fname.c_str()); + } + + if (::fcntl(file_desc, F_RDAHEAD, 0) == -1) { + LOG_WARN("fcntl(F_RDAHEAD, 0) failed for %s (errno=%d: %s)", fname.c_str(), + errno, ::strerror(errno)); } - ctx_map.clear(); } +#endif 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) { @@ -685,32 +892,74 @@ void LinuxAlignedFileReader::open(const std::string &fname) { } #if defined(__APPLE__) || defined(__MACH__) - // macOS has no O_DIRECT. F_NOCACHE is its closest per-file equivalent: it - // asks the kernel to minimize caching for I/O through this descriptor. This - // is advisory rather than a guarantee that every read reaches the device. - // Disable read-ahead as well because DiskAnn performs random reads. - // - // Do not mmap the entire index and call msync(MS_INVALIDATE) here. That does - // not provide a reliable global cache eviction guarantee and makes open time - // and virtual-address usage scale with the size of the index. if (this->file_desc != -1) { - if (::fcntl(this->file_desc, F_NOCACHE, 1) == -1) { + configure_macos_reader(this->file_desc, fname); + } +#endif + + LOG_INFO("Opened file : %s", fname.c_str()); +} + +int LinuxAlignedFileReader::open_from_handle(const std::string &fname, + int source_fd) { + close(); + if (source_fd < 0) { + LOG_ERROR("Cannot capture DiskAnn file from an invalid descriptor"); + return IndexError_InvalidArgument; + } + + int duplicate_fd = -1; + bool has_independent_file_description = false; +#if defined(__linux__) || defined(__linux) + duplicate_fd = reopen_file_descriptor_with_direct_io(source_fd); + if (duplicate_fd < 0) { + const int direct_errno = errno; + duplicate_fd = duplicate_file_descriptor(source_fd); + if (duplicate_fd >= 0) { LOG_WARN( - "fcntl(F_NOCACHE) failed for %s (errno=%d: %s); reads will use " - "the page cache", - fname.c_str(), errno, ::strerror(errno)); - } else { - LOG_INFO("DiskAnn macOS: F_NOCACHE enabled for %s", fname.c_str()); + "opening an independent O_DIRECT descriptor failed for %s " + "(errno=%d: %s); falling back to a buffered duplicate", + fname.c_str(), direct_errno, ::strerror(direct_errno)); } - - if (::fcntl(this->file_desc, F_RDAHEAD, 0) == -1) { - LOG_WARN("fcntl(F_RDAHEAD, 0) failed for %s (errno=%d: %s)", - fname.c_str(), errno, ::strerror(errno)); + } else { + has_independent_file_description = true; + } +#elif defined(__APPLE__) || defined(__MACH__) + duplicate_fd = reopen_macos_file_descriptor(fname, source_fd); + if (duplicate_fd < 0) { + const int reopen_errno = errno; + duplicate_fd = duplicate_file_descriptor(source_fd); + if (duplicate_fd >= 0) { + LOG_WARN( + "opening an independent macOS descriptor failed for %s " + "(errno=%d: %s); falling back to a buffered duplicate", + fname.c_str(), reopen_errno, ::strerror(reopen_errno)); } + } else { + has_independent_file_description = true; } +#else + duplicate_fd = duplicate_file_descriptor(source_fd); #endif + if (duplicate_fd < 0) { + LOG_ERROR( + "Failed to duplicate DiskAnn file descriptor for %s " + "(errno=%d: %s)", + fname.c_str(), errno, ::strerror(errno)); + return IndexError_OpenFile; + } - LOG_INFO("Opened file : %s", fname.c_str()); +#if defined(__APPLE__) || defined(__MACH__) + if (has_independent_file_description) { + configure_macos_reader(duplicate_fd, fname); + } +#else + (void)has_independent_file_description; +#endif + + file_desc = duplicate_fd; + LOG_INFO("Captured open DiskAnn file object: %s", fname.c_str()); + return 0; } void LinuxAlignedFileReader::close() { @@ -755,10 +1004,9 @@ int LinuxAlignedFileReader::submit(PendingBatch &batch, return 0; } - // If this context has no async I/O backend (null/sentinel context or - // explicit pread backend), use synchronous pread. - if (ctx == nullptr || ctx == (IOContext)-1 || - ctx->type == ailego::IOBackendType::kPread) { + // If this context has no async I/O backend (null context or explicit pread + // backend), use synchronous pread. + if (ctx == nullptr || ctx->type == ailego::IOBackendType::kPread) { int pread_ret = execute_io_pread(this->file_desc, read_reqs); if (pread_ret != 0) { return pread_ret; @@ -995,5 +1243,453 @@ int LinuxAlignedFileReader::get_completed( } #endif +#else // Windows + +// Windows uses one file handle and one I/O completion port per IOContext, so a +// context can only dequeue completion packets for requests submitted through +// that context. PendingBatch keeps the expected lengths and completion bitmap +// alive until every request has been harvested. +WindowsAlignedFileReader::~WindowsAlignedFileReader() { + close(); +} + +static bool resolve_windows_file_path(const std::string &fname, + std::wstring &absolute_path) { + const std::wstring wide_fname = ailego::FileHelper::Utf8ToWide(fname); + if (wide_fname.empty()) { + LOG_ERROR("Failed to convert DiskAnn file path from UTF-8: %s", + fname.c_str()); + return false; + } + + const DWORD path_capacity = + ::GetFullPathNameW(wide_fname.c_str(), 0, nullptr, nullptr); + if (path_capacity == 0) { + LOG_ERROR("Failed to resolve absolute DiskAnn file path: %s (error=%lu)", + fname.c_str(), ::GetLastError()); + return false; + } + absolute_path.assign(path_capacity, L'\0'); + const DWORD path_length = ::GetFullPathNameW( + wide_fname.c_str(), path_capacity, absolute_path.data(), nullptr); + if (path_length == 0 || path_length >= path_capacity) { + LOG_ERROR("Failed to resolve absolute DiskAnn file path: %s (error=%lu)", + fname.c_str(), ::GetLastError()); + absolute_path.clear(); + return false; + } + absolute_path.resize(path_length); + return true; +} + +void WindowsAlignedFileReader::open(const std::string &fname) { + close(); + std::wstring absolute_path; + if (!resolve_windows_file_path(fname, absolute_path)) { + return; + } + + HANDLE stable_file_handle = ::CreateFileW( + absolute_path.c_str(), GENERIC_READ, kDiskAnnFileShareMode, nullptr, + OPEN_EXISTING, FILE_ATTRIBUTE_READONLY | kDiskAnnStableHandleFlags, + nullptr); + if (stable_file_handle == INVALID_HANDLE_VALUE) { + LOG_ERROR("Failed to open file: %s (error=%lu)", fname.c_str(), + ::GetLastError()); + return; + } + const uint64_t file_identity = next_windows_file_identity(); + if (file_identity == 0) { + ::CloseHandle(stable_file_handle); + LOG_ERROR("Exhausted DiskAnn Windows file identities"); + return; + } + stable_file_handle_ = stable_file_handle; + file_path_ = std::move(absolute_path); + file_identity_ = file_identity; + LOG_INFO("Opened file: %s", fname.c_str()); +} + +int WindowsAlignedFileReader::open_from_handle(const std::string &fname, + HANDLE source_handle) { + close(); + if (source_handle == INVALID_HANDLE_VALUE || source_handle == nullptr) { + LOG_ERROR("Cannot capture DiskAnn file from an invalid source handle"); + return IndexError_InvalidArgument; + } + + std::wstring absolute_path; + if (!resolve_windows_file_path(fname, absolute_path)) { + return IndexError_InvalidArgument; + } + + // ReOpenFile refers to the same underlying file object even if fname has + // already been renamed or replaced. Keep this stable handle unbuffered too: + // an ordinary buffered handle beside the private unbuffered IOCP handles can + // severely reduce random-read throughput. + HANDLE stable_file_handle = + ::ReOpenFile(source_handle, GENERIC_READ, kDiskAnnFileShareMode, + kDiskAnnStableHandleFlags); + if (stable_file_handle == INVALID_HANDLE_VALUE) { + LOG_ERROR("Failed to capture DiskAnn file object (error=%lu)", + ::GetLastError()); + return IndexError_Runtime; + } + const uint64_t file_identity = next_windows_file_identity(); + if (file_identity == 0) { + ::CloseHandle(stable_file_handle); + LOG_ERROR("Exhausted DiskAnn Windows file identities"); + return IndexError_Runtime; + } + + stable_file_handle_ = stable_file_handle; + file_path_ = std::move(absolute_path); + file_identity_ = file_identity; + LOG_INFO("Captured open DiskAnn file object: %s", fname.c_str()); + return 0; +} + +void WindowsAlignedFileReader::close() { + if (stable_file_handle_ != INVALID_HANDLE_VALUE) { + ::CloseHandle(stable_file_handle_); + stable_file_handle_ = INVALID_HANDLE_VALUE; + } + file_path_.clear(); + file_identity_ = 0; +} + +int WindowsAlignedFileReader::prepare_io_ctx(IOContext &ctx) { + if (ctx == nullptr || + ctx->type != ailego::IOBackendType::kWindowsOverlapped) { + LOG_ERROR("Attempt to prepare an invalid Windows I/O context"); + return IndexError_Runtime; + } + if (file_path_.empty() || stable_file_handle_ == INVALID_HANDLE_VALUE) { + LOG_ERROR("Attempt to read before opening a DiskAnn file"); + return IndexError_Runtime; + } + if (ctx->file_handle != INVALID_HANDLE_VALUE && + ctx->completion_port != nullptr && ctx->file_path == file_path_ && + ctx->file_identity == file_identity_) { + return 0; + } + if (ctx->outstanding_count != 0) { + LOG_ERROR("Cannot replace a Windows I/O context with requests in flight"); + return IndexError_Runtime; + } + + close_windows_io_handles(ctx); + // Derive each private IOCP handle from the file object captured by open(). + // Reopening the path here could bind a lazy context to a replacement index + // while the indexer still holds metadata for the original one. + ctx->file_handle = + ::ReOpenFile(stable_file_handle_, GENERIC_READ, kDiskAnnFileShareMode, + FILE_FLAG_NO_BUFFERING | FILE_FLAG_OVERLAPPED); + if (ctx->file_handle == INVALID_HANDLE_VALUE) { + LOG_ERROR("Failed to reopen DiskAnn file object for IOCP (error=%lu)", + ::GetLastError()); + return IndexError_Runtime; + } + + ctx->completion_port = ::CreateIoCompletionPort( + ctx->file_handle, nullptr, reinterpret_cast(ctx), + kDiskAnnIoCompletionConcurrency); + if (ctx->completion_port == nullptr) { + LOG_ERROR("CreateIoCompletionPort failed (error=%lu)", ::GetLastError()); + close_windows_io_handles(ctx); + return IndexError_Runtime; + } + if (!::SetFileCompletionNotificationModes(ctx->file_handle, + FILE_SKIP_SET_EVENT_ON_HANDLE)) { + LOG_WARN("SetFileCompletionNotificationModes failed (error=%lu)", + ::GetLastError()); + } + try { + ctx->file_path = file_path_; + ctx->file_identity = file_identity_; + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to store the Windows DiskAnn file path"); + close_windows_io_handles(ctx); + return IndexError_NoMemory; + } + return 0; +} + +void WindowsAlignedFileReader::reset_io_ctx(IOContext &ctx) { + close_windows_io_handles(ctx); +} + +void WindowsAlignedFileReader::release_io_ctx(IOContext &ctx) { + close_windows_io_handles(ctx); +} + +static int validate_windows_read_requests( + const std::vector &read_reqs) { + constexpr uint64_t kSectorLen = 4096; + for (size_t i = 0; i < read_reqs.size(); ++i) { + const AlignedRead &req = read_reqs[i]; + if (req.buf == nullptr || + reinterpret_cast(req.buf) % kSectorLen != 0 || + req.offset % kSectorLen != 0 || req.len % kSectorLen != 0) { + LOG_ERROR( + "Invalid unbuffered read request %zu: buffer=%p, offset=%llu, " + "len=%llu; all values must be aligned to %llu bytes", + i, req.buf, static_cast(req.offset), + static_cast(req.len), + static_cast(kSectorLen)); + return IndexError_InvalidArgument; + } + if (req.len > (std::numeric_limits::max)()) { + LOG_ERROR("Windows read request %zu is too large: %llu bytes", i, + static_cast(req.len)); + return IndexError_InvalidArgument; + } + } + return 0; +} + +int WindowsAlignedFileReader::read(std::vector &read_reqs, + IOContext &ctx, bool async) { + if (async) { + LOG_WARN( + "read() waits for completion; use submit()/get_completed() for " + "asynchronous Windows I/O"); + } + int ret = validate_windows_read_requests(read_reqs); + if (ret != 0) { + return ret; + } + + std::vector completed; + try { + completed.reserve(MAX_IO_DEPTH); + } catch (const std::bad_alloc &) { + return IndexError_NoMemory; + } + + for (size_t start = 0; start < read_reqs.size(); start += MAX_IO_DEPTH) { + const size_t count = + std::min(read_reqs.size() - start, MAX_IO_DEPTH); + std::vector requests(read_reqs.begin() + start, + read_reqs.begin() + start + count); + PendingBatch batch; + ret = submit(batch, requests, ctx); + if (ret != 0) { + return ret; + } + + while (batch.n_reaped < batch.n_submitted) { + ret = get_completed(batch, ctx, 1, completed); + if (ret < 0) { + return ret; + } + } + } + return 0; +} + +int WindowsAlignedFileReader::submit(PendingBatch &batch, + std::vector &read_reqs, + IOContext &ctx) { + batch.n_submitted = 0; + batch.n_reaped = 0; + batch.used_pread = false; + batch.expected_lengths.clear(); + batch.completed.clear(); + batch.generation = 0; + + if (read_reqs.empty()) { + return 0; + } + if (read_reqs.size() > MAX_IO_DEPTH) { + LOG_ERROR("Windows IOCP batch has %zu requests; maximum is %u", + read_reqs.size(), static_cast(MAX_IO_DEPTH)); + return IndexError_InvalidArgument; + } + int ret = validate_windows_read_requests(read_reqs); + if (ret != 0) { + return ret; + } + ret = prepare_io_ctx(ctx); + if (ret != 0) { + return ret; + } + if (ctx->outstanding_count != 0) { + LOG_ERROR("Windows I/O context already has an active batch"); + return IndexError_Runtime; + } + + ++ctx->generation; + if (ctx->generation == 0) { + ++ctx->generation; + } + batch.generation = ctx->generation; + try { + batch.expected_lengths.reserve(read_reqs.size()); + batch.completed.assign(read_reqs.size(), 0); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate Windows IOCP batch metadata"); + batch.expected_lengths.clear(); + batch.completed.clear(); + batch.generation = 0; + return IndexError_NoMemory; + } + + uint32_t issued_count = 0; + ctx->active_requests.fill(0); + for (size_t i = 0; i < read_reqs.size(); ++i) { + ctx->reqs[i] = OVERLAPPED{}; + + const AlignedRead &req = read_reqs[i]; + OVERLAPPED &request = ctx->reqs[i]; + request.Offset = static_cast(req.offset & 0xffffffffULL); + request.OffsetHigh = static_cast(req.offset >> 32); + + BOOL queued = ::ReadFile(ctx->file_handle, req.buf, + static_cast(req.len), nullptr, &request); + if (!queued && ::GetLastError() != ERROR_IO_PENDING) { + LOG_ERROR("Error queuing IOCP read %zu (error=%lu)", i, ::GetLastError()); + ctx->outstanding_count = issued_count; + reset_io_ctx(ctx); + batch.expected_lengths.clear(); + batch.completed.clear(); + batch.generation = 0; + return IndexError_Runtime; + } + + batch.expected_lengths.push_back(req.len); + ctx->active_requests[i] = 1; + ++issued_count; + ctx->outstanding_count = issued_count; + } + + batch.n_submitted = issued_count; + return 0; +} + +int WindowsAlignedFileReader::get_completed( + PendingBatch &batch, IOContext &ctx, int min_completed, + std::vector &completed_indices) { + completed_indices.clear(); + if (batch.n_reaped >= batch.n_submitted) { + return 0; + } + if (ctx == nullptr || ctx->completion_port == nullptr || + batch.generation == 0 || batch.generation != ctx->generation || + batch.expected_lengths.size() != batch.n_submitted || + batch.completed.size() != batch.n_submitted || + ctx->outstanding_count != batch.n_submitted - batch.n_reaped) { + LOG_ERROR("Invalid or stale Windows IOCP batch"); + reset_io_ctx(ctx); + batch.n_reaped = batch.n_submitted; + return IndexError_Runtime; + } + + if (completed_indices.capacity() < ctx->outstanding_count) { + try { + completed_indices.reserve(ctx->outstanding_count); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate Windows IOCP completion metadata"); + reset_io_ctx(ctx); + batch.n_reaped = batch.n_submitted; + return IndexError_NoMemory; + } + } + + const uint32_t remaining = batch.n_submitted - batch.n_reaped; + const uint32_t target = std::min( + remaining, static_cast(std::max(min_completed, 1))); + + while (completed_indices.size() < target) { + OVERLAPPED_ENTRY entries[MAX_IO_DEPTH]{}; + ULONG removed = 0; + const ULONG max_entries = static_cast(std::min( + ctx->outstanding_count, static_cast(MAX_IO_DEPTH))); + BOOL dequeued = ::GetQueuedCompletionStatusEx( + ctx->completion_port, entries, max_entries, &removed, INFINITE, FALSE); + if (!dequeued || removed == 0) { + LOG_ERROR("GetQueuedCompletionStatusEx failed (error=%lu)", + ::GetLastError()); + reset_io_ctx(ctx); + batch.n_reaped = batch.n_submitted; + completed_indices.clear(); + return IndexError_Runtime; + } + + bool completion_error = false; + for (ULONG i = 0; i < removed; ++i) { + if (entries[i].lpCompletionKey != reinterpret_cast(ctx)) { + LOG_ERROR("IOCP returned a completion for a different context"); + completion_error = true; + continue; + } + + const uintptr_t address = + reinterpret_cast(entries[i].lpOverlapped); + const uintptr_t begin = reinterpret_cast(ctx->reqs.data()); + const uintptr_t span = + static_cast(batch.n_submitted) * sizeof(OVERLAPPED); + if (address < begin || address >= begin + span || + (address - begin) % sizeof(OVERLAPPED) != 0) { + LOG_ERROR("IOCP returned an unknown OVERLAPPED request"); + completion_error = true; + continue; + } + + const uint32_t index = + static_cast((address - begin) / sizeof(OVERLAPPED)); + if (batch.completed[index] != 0) { + LOG_ERROR("IOCP returned duplicate completion for request %u", index); + completion_error = true; + continue; + } + + if (ctx->active_requests[index] == 0 || ctx->outstanding_count == 0) { + LOG_ERROR("IOCP returned an inactive or excess completion"); + completion_error = true; + continue; + } + + DWORD bytes_transferred = 0; + bool terminal = true; + if (!::GetOverlappedResult(ctx->file_handle, &ctx->reqs[index], + &bytes_transferred, FALSE)) { + DWORD error = ::GetLastError(); + terminal = error != ERROR_IO_INCOMPLETE; + LOG_ERROR("IOCP read %u failed (error=%lu)", index, error); + completion_error = true; + } else if (static_cast(bytes_transferred) != + batch.expected_lengths[index]) { + LOG_ERROR( + "IOCP read %u completed with %lu bytes, expected %llu", index, + static_cast(bytes_transferred), + static_cast(batch.expected_lengths[index])); + completion_error = true; + } + + if (terminal) { + ctx->active_requests[index] = 0; + --ctx->outstanding_count; + batch.completed[index] = 1; + ++batch.n_reaped; + } + if (!completion_error) { + completed_indices.push_back(index); + } + } + + if (completion_error) { + reset_io_ctx(ctx); + batch.n_reaped = batch.n_submitted; + completed_indices.clear(); + return IndexError_Runtime; + } + } + + return static_cast(completed_indices.size()); +} + +#endif // Windows/POSIX reader implementation + } // namespace core } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_file_reader.h b/src/core/algorithm/diskann/diskann_file_reader.h index 809b68c7f..d9589050a 100644 --- a/src/core/algorithm/diskann/diskann_file_reader.h +++ b/src/core/algorithm/diskann/diskann_file_reader.h @@ -16,16 +16,34 @@ #define MAX_IO_DEPTH 128 #include +#include #if (defined(__linux) || defined(__linux__)) #include // raw-syscall io_uring wrapper (IoUringRing) #include // dlopen-based libaio wrapper +#elif defined(_WIN32) || defined(_WIN64) +#ifndef NOMINMAX +#define NOMINMAX #endif +#ifndef _WIN32_WINNT +#define _WIN32_WINNT 0x0600 +#endif +#include +// Do not leak Win32's function-like aliases into headers included below. +// They otherwise rewrite qualified names such as FileHelper::DeleteFile(). +#ifdef DeleteFile +#undef DeleteFile +#endif +#ifdef RemoveDirectory +#undef RemoveDirectory +#endif +#endif + +#if !defined(_WIN32) && !defined(_WIN64) #include -#include -#include -#include +#endif +#include #include #include #include @@ -42,16 +60,25 @@ namespace core { // // macOS uses a real context with type kPread so that the active backend can be // inspected and reported consistently instead of using an opaque placeholder. -// IOContext is a pointer to IoBackend, which preserves the existing -// sentinel conventions: nullptr means uninitialised and (IOContext)-1 is -// the invalid-handle sentinel returned by get_ctx() for unregistered -// threads. +// Windows stores a private file handle, completion port, and stable OVERLAPPED +// request slots in each I/O context. Keeping completion ports private prevents +// one context from consuming another context's completions. +// IOContext is a pointer to IoBackend; nullptr means uninitialised. struct IoBackend { ailego::IOBackendType type{ailego::IOBackendType::kPread}; #if (defined(__linux) || defined(__linux__)) IoUringRing ring{}; io_context_t aio_ctx{nullptr}; +#elif defined(_WIN32) || defined(_WIN64) + std::array reqs{}; + std::array active_requests{}; + HANDLE file_handle{INVALID_HANDLE_VALUE}; + HANDLE completion_port{nullptr}; + std::wstring file_path; + uint64_t file_identity{0}; + uint32_t outstanding_count{0}; + uint64_t generation{0}; #endif }; @@ -61,7 +88,8 @@ int setup_io_ctx(IOContext &ctx); int destroy_io_ctx(IOContext &ctx); // Log the current DiskAnn I/O backend (io_uring, libaio, or pread). Probes the -// backend on first call. No-op outside Linux and macOS. +// backend on first call. No-op outside Linux and macOS; Android and iOS always +// use synchronous pread. void log_diskann_io_backend(); struct AlignedRead { @@ -86,6 +114,10 @@ struct PendingBatch { #if (defined(__linux) || defined(__linux__)) std::vector cbs; std::vector cb_ptrs; +#elif defined(_WIN32) || defined(_WIN64) + std::vector expected_lengths; + std::vector completed; + uint64_t generation{0}; #endif uint32_t n_submitted{0}; uint32_t n_reaped{0}; @@ -93,19 +125,9 @@ struct PendingBatch { }; class AlignedFileReader { - protected: - std::map ctx_map; - std::mutex ctx_mut; - public: - virtual IOContext &get_ctx() = 0; - virtual ~AlignedFileReader() {} - virtual void register_thread() = 0; - virtual void deregister_thread() = 0; - virtual void deregister_all_threads() = 0; - virtual void open(const std::string &fname) = 0; virtual void close() = 0; @@ -118,39 +140,85 @@ class AlignedFileReader { virtual int get_completed(PendingBatch &batch, IOContext &ctx, int min_completed, std::vector &completed_indices) = 0; + + // Release any lazy per-context file resources at an operation boundary. + // POSIX backends keep their process-local queue resources for reuse; Windows + // overrides this to close private file and completion-port handles. + virtual void release_io_ctx(IOContext &ctx) = 0; }; -// Reader implementation used on all supported platforms. Linux selects -// io_uring, libaio, or pread. macOS ARM64 uses synchronous pread. +// POSIX reader implementation. Linux selects io_uring, libaio, or pread; +// macOS ARM64 uses synchronous pread. +#if !defined(_WIN32) && !defined(_WIN64) class LinuxAlignedFileReader : public AlignedFileReader { private: int file_desc; - IOContext bad_ctx = (IOContext)-1; - public: LinuxAlignedFileReader(); LinuxAlignedFileReader(int file_desc); - ~LinuxAlignedFileReader(); + LinuxAlignedFileReader(const LinuxAlignedFileReader &) = delete; + LinuxAlignedFileReader &operator=(const LinuxAlignedFileReader &) = delete; + ~LinuxAlignedFileReader() override; public: - IOContext &get_ctx(); - - void register_thread(); - void deregister_thread(); - void deregister_all_threads(); - void open(const std::string &fname); - void close(); + void open(const std::string &fname) override; + // Duplicate an already-open descriptor so metadata and graph reads stay on + // the same file object even if fname is atomically replaced. + int open_from_handle(const std::string &fname, int source_fd); + void close() override; int read(std::vector &read_reqs, IOContext &ctx, - bool async = false); + bool async = false) override; int submit(PendingBatch &batch, std::vector &read_reqs, - IOContext &ctx); + IOContext &ctx) override; int get_completed(PendingBatch &batch, IOContext &ctx, int min_completed, - std::vector &completed_indices); + std::vector &completed_indices) override; + void release_io_ctx(IOContext & /*ctx*/) override {} }; +#else +class WindowsAlignedFileReader : public AlignedFileReader { + private: + friend class WindowsAlignedFileReaderTestPeer; + + std::wstring file_path_; + HANDLE stable_file_handle_{INVALID_HANDLE_VALUE}; + uint64_t file_identity_{0}; + + int prepare_io_ctx(IOContext &ctx); + void reset_io_ctx(IOContext &ctx); + + public: + WindowsAlignedFileReader() = default; + WindowsAlignedFileReader(const WindowsAlignedFileReader &) = delete; + WindowsAlignedFileReader &operator=(const WindowsAlignedFileReader &) = + delete; + ~WindowsAlignedFileReader() override; + + void open(const std::string &fname) override; + // Capture the same file object as an already-open buffered handle. This is + // used while loading an index so metadata and later graph reads cannot come + // from different files if fname is atomically replaced between the two. + int open_from_handle(const std::string &fname, HANDLE source_handle); + void close() override; + + int read(std::vector &read_reqs, IOContext &ctx, + bool async = false) override; + int submit(PendingBatch &batch, std::vector &read_reqs, + IOContext &ctx) override; + int get_completed(PendingBatch &batch, IOContext &ctx, int min_completed, + std::vector &completed_indices) override; + void release_io_ctx(IOContext &ctx) override; +}; +#endif + +#if defined(_WIN32) || defined(_WIN64) +using PlatformAlignedFileReader = WindowsAlignedFileReader; +#else +using PlatformAlignedFileReader = LinuxAlignedFileReader; +#endif } // namespace core } // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_holder.h b/src/core/algorithm/diskann/diskann_holder.h deleted file mode 100644 index beb422ba2..000000000 --- a/src/core/algorithm/diskann/diskann_holder.h +++ /dev/null @@ -1,325 +0,0 @@ -// 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. -#pragma once - -#include -#include -#include "diskann_entity.h" - -namespace zvec { -namespace core { - -struct DiskAnnIndexHolderMeta { - uint32_t element_size_; - uint32_t key_size_; - uint32_t sector_size_; - uint32_t doc_cnt_; - uint8_t reserve_[]; -}; - -class DiskAnnIndexHolder : public IndexHolder { - public: - typedef std::shared_ptr Pointer; - - public: - enum Status { STATUS_UNINITED = 0, STATUS_WRITE = 1, STATUS_READ = 2 }; - - public: - static constexpr uint32_t kDataSectorSize = 128 * 1024; - static constexpr uint32_t kMetaSectorSize = 4096; - - public: - inline static uint32_t get_sector_id(uint32_t id, uint32_t sector_vec_num) { - return id / sector_vec_num; - } - - inline static uint32_t get_sector_offset(uint32_t id, uint32_t sector_vec_num, - uint32_t data_size) { - return (id % sector_vec_num) * data_size; - } - - public: - /*! Random Access Index Holder Iterator - */ - class Iterator : public IndexHolder::Iterator { - public: - //! Index Holder Iterator Pointer - typedef std::unique_ptr Pointer; - - //! Constructor - Iterator(DiskAnnIndexHolder *owner) - : holder_(owner), sector_id_{0}, sector_offset_{0} { - path_ = holder_->path(); - data_size_ = holder_->data_size(); - data_sector_size_ = holder_->data_sector_size(); - meta_sector_size_ = holder_->meta_sector_size(); - - sector_buffer_.resize(data_sector_size_); - - sector_vec_num_ = data_sector_size_ / data_size_; - } - - //! Destructor - virtual ~Iterator(void) { - if (file_.is_open()) { - file_.close(); - } - } - - int init() { - file_.open(path_, std::ios::in); - if (!file_.is_open()) { - LOG_ERROR("file can not create, %s", path_.c_str()); - return IndexError_OpenFile; - } - - file_.seekg(meta_sector_size_); - - read_sector(); - - return 0; - } - - //! Retrieve pointer of data - const void *data(void) const override { - const uint8_t *data_ptr = - reinterpret_cast(sector_buffer_.data()); - return data_ptr + sector_offset_ + sizeof(diskann_key_t); - } - - //! Test if the iterator is valid - bool is_valid(void) const override { - return id_ < holder_->count(); - } - - //! Retrieve primary key - uint64_t key(void) const override { - const uint8_t *data_ptr = - reinterpret_cast(sector_buffer_.data()); - uint64_t key = - *reinterpret_cast(data_ptr + sector_offset_); - - return key; - } - - //! Next iterator - void next(void) override { - ++id_; - - uint32_t sector_id = get_sector_id(id_, sector_vec_num_); - if (sector_id > sector_id_) { - file_.seekg(sector_id * data_sector_size_ + meta_sector_size_); - read_sector(); - sector_id_ = sector_id; - } - - sector_offset_ = get_sector_offset(id_, sector_vec_num_, data_size_); - } - - int read_sector() { - file_.read(&((sector_buffer_)[0]), data_sector_size_); - if (!file_) { - LOG_ERROR("Failed to read sector from file: %s", path_.c_str()); - return IndexError_ReadData; - } - - return 0; - } - - private: - //! Members - DiskAnnIndexHolder *holder_{nullptr}; - std::string path_; - std::ifstream file_; - uint32_t sector_id_{0}; - std::string sector_buffer_; - uint32_t id_{0}; - uint32_t data_size_{0}; - uint32_t sector_offset_{0}; - uint32_t data_sector_size_{0}; - uint32_t meta_sector_size_{0}; - uint32_t sector_vec_num_{0}; - }; - - public: - DiskAnnIndexHolder(IndexMeta &meta, std::string &path) { - path_ = path; - - data_size_ = meta.element_size() + sizeof(diskann_key_t); - dimension_ = meta.dimension(); - type_ = meta.data_type(); - - element_size_ = meta.element_size(); - sector_vec_num_ = data_sector_size_ / data_size_; - padding_size_ = data_sector_size_ - sector_vec_num_ * data_size_; - sector_buffer_.resize(data_sector_size_); - sector_internal_id_ = 0; - } - - ~DiskAnnIndexHolder() override { - if (file_.is_open()) { - file_.close(); - } - } - - //! Init - int init() { - file_.open(path_, std::ios::out | std::ios::trunc); - - if (!file_.is_open()) { - LOG_ERROR("file can not create, %s", path_.c_str()); - return IndexError_OpenFile; - } - - DiskAnnIndexHolderMeta holder_meta; - holder_meta.element_size_ = element_size_; - holder_meta.key_size_ = sizeof(diskann_key_t); - holder_meta.sector_size_ = data_sector_size_; - - std::vector empty_sector; - empty_sector.resize(meta_sector_size_); - - std::memset(&(empty_sector[0]), 0, meta_sector_size_); - std::memcpy(&(empty_sector[0]), &holder_meta, - sizeof(DiskAnnIndexHolderMeta)); - - file_.write(reinterpret_cast(&(empty_sector[0])), - meta_sector_size_); - if (!file_) { - LOG_ERROR("Failed to write meta sector to file: %s", path_.c_str()); - return IndexError_WriteData; - } - - status_ = STATUS_WRITE; - - return 0; - } - - int close() { - if (sector_internal_id_ != 0) { - file_.write(reinterpret_cast(&(sector_buffer_[0])), - data_sector_size_); - } - - file_.close(); - - return 0; - } - - //! Retrieve count of elements in holder (-1 indicates unknown) - size_t count(void) const override { - return count_; - } - - //! Retrieve dimension - size_t dimension(void) const override { - return dimension_; - } - - //! Retrieve type information - IndexMeta::DataType data_type(void) const override { - return type_; - } - - //! Retrieve element size in bytes - size_t element_size(void) const override { - return element_size_; - } - - //! Retrieve if it can multi-pass - bool multipass(void) const override { - return true; - } - - //! Create a new iterator - IndexHolder::Iterator::Pointer create_iterator(void) override { - auto pointer = std::make_unique(this); - - if (pointer->init() != 0) { - return nullptr; - } - - return pointer; - } - - int emplace(uint64_t pkey, const void *vec) { - if (status_ != STATUS_WRITE) { - return IndexError_NoReady; - } - - uint8_t *data_ptr = reinterpret_cast(&(sector_buffer_[0])) + - sector_internal_id_ * data_size_; - std::memcpy(data_ptr, &pkey, sizeof(diskann_key_t)); - std::memcpy(data_ptr + sizeof(diskann_key_t), vec, element_size_); - - sector_internal_id_++; - if (sector_internal_id_ >= sector_vec_num_) { - std::memset(data_ptr + data_size_, 0, padding_size_); - - file_.write(reinterpret_cast(&(sector_buffer_[0])), - data_sector_size_); - - sector_internal_id_ = 0; - sector_id_++; - } - - count_++; - - return 0; - } - - uint32_t data_sector_size() { - return data_sector_size_; - } - - uint32_t meta_sector_size() { - return meta_sector_size_; - } - - uint32_t data_size() { - return data_size_; - } - - uint32_t *mutable_sector_id() { - return §or_id_; - } - - uint32_t sector_id() { - return sector_id_; - } - - std::string &path() { - return path_; - } - - private: - std::string path_; - std::ofstream file_; - uint32_t element_size_{0}; - uint32_t dimension_{0}; - IndexMeta::DataType type_{IndexMeta::DataType::DT_UNDEFINED}; - uint32_t sector_vec_num_{0}; - uint32_t data_size_{0}; - uint32_t padding_size_{0}; - uint32_t meta_sector_size_{DiskAnnIndexHolder::kMetaSectorSize}; - uint32_t data_sector_size_{DiskAnnIndexHolder::kDataSectorSize}; - std::string sector_buffer_; - uint32_t sector_internal_id_{0}; - uint32_t sector_id_{0}; - uint32_t count_{0}; - uint32_t status_{STATUS_UNINITED}; -}; - -} // namespace core -} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_index_provider.h b/src/core/algorithm/diskann/diskann_index_provider.h index 1fd8754bf..fb481b89c 100644 --- a/src/core/algorithm/diskann/diskann_index_provider.h +++ b/src/core/algorithm/diskann/diskann_index_provider.h @@ -13,34 +13,77 @@ // limitations under the License. #pragma once +#include +#include +#include +#include #include #include #include -#include "diskann_entity.h" +#include "diskann_context.h" +#include "diskann_indexer.h" namespace zvec { namespace core { -//! IndexProvider implementation backed by a DiskAnn entity. +class DiskAnnProviderTestPeer; + +//! IndexProvider implementation backed by a DiskAnn indexer. //! //! Used by ``MixedStreamerReducer`` during segment merge: the reducer needs //! to walk every vector held by a source DiskAnn streamer and feed it into -//! the merge target. Vectors are read on demand from the entity's on-disk -//! vector segment via ``DiskAnnEntity::get_vector(id)``. +//! the merge target. Vectors are read on demand through the same aligned file +//! reader used by DiskAnn search. The provider owns the indexer, its in-memory +//! entity and independent I/O contexts, so it remains valid after its source +//! streamer is closed. class DiskAnnIndexProvider : public IndexProvider { + friend class DiskAnnProviderTestPeer; + + private: + struct ResultBufferOwner {}; + + struct ThreadResultBuffer { + explicit ThreadResultBuffer( + const std::shared_ptr &buffer_owner) + : owner(buffer_owner) {} + + std::weak_ptr owner; + std::string data; + }; + public: DiskAnnIndexProvider(const IndexMeta &meta, + const IndexMetric::Pointer &measure, const DiskAnnEntity::Pointer &entity, + const DiskAnnIndexer::Pointer &indexer, const std::string &owner) - : meta_(meta), entity_(entity), owner_class_(owner) {} + : meta_(meta), + measure_(measure), + entity_(entity), + indexer_(indexer), + owner_class_(owner) { + try { + result_buffer_owner_ = std::make_shared(); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate DiskAnn provider result-buffer owner"); + } + } DiskAnnIndexProvider(const DiskAnnIndexProvider &) = delete; DiskAnnIndexProvider &operator=(const DiskAnnIndexProvider &) = delete; public: IndexProvider::Iterator::Pointer create_iterator() override { - return IndexProvider::Iterator::Pointer(new (std::nothrow) - Iterator(entity_)); + std::unique_ptr iterator( + new (std::nothrow) Iterator(meta_, measure_, entity_, indexer_)); + if (!iterator || !iterator->ready()) { + return nullptr; + } + return IndexProvider::Iterator::Pointer(iterator.release()); + } + + bool ready() const { + return measure_ && entity_ && indexer_ && result_buffer_owner_; } size_t count(void) const override { @@ -60,11 +103,38 @@ class DiskAnnIndexProvider : public IndexProvider { } const void *get_vector(uint64_t key) const override { - diskann_id_t id = entity_->get_id(static_cast(key)); + if (!ready()) { + return nullptr; + } + + const diskann_id_t id = indexer_->get_id(static_cast(key)); if (id == kInvalidId) { return nullptr; } - return entity_->get_vector(id); + + // Serialize the heavyweight I/O context, but keep only the returned bytes + // in thread-local storage. Buffers are isolated by provider lifetime, so a + // fetch through another provider or on another thread cannot invalidate + // this pointer. Expired providers are pruned lazily without retaining any + // per-thread file/context resources. The returned pointer is valid until + // this thread's next fetch through this provider, provider destruction, or + // thread exit. + try { + std::string &vector_buffer = thread_result_buffer(result_buffer_owner_); + std::lock_guard lock(fetch_mutex_); + if (!fetch_context_) { + fetch_context_ = + DiskAnnContext::create_fetch_context(meta_, measure_, entity_); + } + if (!fetch_context_ || + indexer_->get_vector(id, fetch_context_, vector_buffer) != 0) { + return nullptr; + } + return vector_buffer.data(); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate DiskAnn provider vector buffer"); + return nullptr; + } } const std::string &owner_class(void) const override { @@ -72,15 +142,65 @@ class DiskAnnIndexProvider : public IndexProvider { } private: + static std::string &thread_result_buffer( + const std::shared_ptr &owner) { + // A list keeps every live provider's string object stable when another + // provider first fetches on this thread. Only the vector bytes are kept in + // TLS; heavyweight DiskAnnContext and file handles remain provider-owned. + static thread_local std::list buffers; + for (auto it = buffers.begin(); it != buffers.end();) { + std::shared_ptr entry_owner = it->owner.lock(); + if (!entry_owner) { + it = buffers.erase(it); + continue; + } + if (entry_owner.get() == owner.get()) { + return it->data; + } + ++it; + } + + buffers.emplace_back(owner); + return buffers.back().data; + } + class Iterator : public IndexProvider::Iterator { public: - explicit Iterator(const DiskAnnEntity::Pointer &entity) - : entity_(entity), cur_id_(0U) { + Iterator(const IndexMeta &meta, const IndexMetric::Pointer &measure, + const DiskAnnEntity::Pointer &entity, + const DiskAnnIndexer::Pointer &indexer) + : meta_(meta), + measure_(measure), + entity_(entity), + indexer_(indexer), + cur_id_(0U) { cur_id_ = next_valid_id(0U); } + bool ready() const { + return meta_.element_size() > 0 && measure_ && entity_ && indexer_; + } + const void *data(void) const override { - return entity_->get_vector(cur_id_); + if (!is_valid() || !ready()) { + return nullptr; + } + if (!data_loaded_) { + // Context setup owns aligned I/O scratch space and platform resources; + // iterators that are only inspected should not pay that cost. + if (!context_) { + context_ = + DiskAnnContext::create_fetch_context(meta_, measure_, entity_); + } + if (!context_) { + return nullptr; + } + if (indexer_->get_vector(cur_id_, context_, vector_buffer_) != 0) { + return nullptr; + } + data_loaded_ = true; + } + return vector_buffer_.data(); } bool is_valid(void) const override { @@ -93,9 +213,13 @@ class DiskAnnIndexProvider : public IndexProvider { void next(void) override { cur_id_ = next_valid_id(cur_id_ + 1); + data_loaded_ = false; + vector_buffer_.clear(); } private: + friend class DiskAnnProviderTestPeer; + //! Skip ids that map to ``kInvalidKey`` (deleted / never populated slots). diskann_id_t next_valid_id(diskann_id_t start_id) const { const auto total = static_cast(entity_->doc_cnt()); @@ -107,13 +231,24 @@ class DiskAnnIndexProvider : public IndexProvider { return total; } + IndexMeta meta_; + IndexMetric::Pointer measure_; DiskAnnEntity::Pointer entity_; + DiskAnnIndexer::Pointer indexer_; + mutable IndexContext::Pointer context_; + mutable std::string vector_buffer_; + mutable bool data_loaded_{false}; diskann_id_t cur_id_; }; IndexMeta meta_; + IndexMetric::Pointer measure_; DiskAnnEntity::Pointer entity_; + DiskAnnIndexer::Pointer indexer_; std::string owner_class_; + std::shared_ptr result_buffer_owner_; + mutable std::mutex fetch_mutex_; + mutable IndexContext::Pointer fetch_context_; }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_indexer.cc b/src/core/algorithm/diskann/diskann_indexer.cc index 98ffcad2e..c45d82181 100644 --- a/src/core/algorithm/diskann/diskann_indexer.cc +++ b/src/core/algorithm/diskann/diskann_indexer.cc @@ -14,15 +14,66 @@ #include "diskann_indexer.h" #include +#include +#include #include +#include #include +#include #include #include #include +#include namespace zvec { namespace core { +namespace { + +// DiskAnnContext instances are pooled above the indexer and can therefore +// outlive the reader that prepared their lazy Windows file handle. Keep that +// handle across all I/O batches in one logical operation, then release it on +// every return path (including exceptions). POSIX readers intentionally keep +// their backend queue resources for reuse. +class IOContextReleaseGuard { + public: + IOContextReleaseGuard(AlignedFileReader &reader, IOContext &ctx, + bool enabled = true) + : reader_(reader), ctx_(ctx), enabled_(enabled) {} + + ~IOContextReleaseGuard() { + if (enabled_) { + reader_.release_io_ctx(ctx_); + } + } + + IOContextReleaseGuard(const IOContextReleaseGuard &) = delete; + IOContextReleaseGuard &operator=(const IOContextReleaseGuard &) = delete; + + private: + AlignedFileReader &reader_; + IOContext &ctx_; + bool enabled_; +}; + +bool checked_multiply_u64(uint64_t lhs, uint64_t rhs, uint64_t *result) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return false; + } + *result = lhs * rhs; + return true; +} + +bool checked_add_u64(uint64_t lhs, uint64_t rhs, uint64_t *result) { + if (rhs > std::numeric_limits::max() - lhs) { + return false; + } + *result = lhs + rhs; + return true; +} + +} // namespace + DiskAnnIndexer::DiskAnnIndexer(const IndexMeta &meta) { meta_ = meta; } @@ -30,61 +81,222 @@ DiskAnnIndexer::DiskAnnIndexer(const IndexMeta &meta) { DiskAnnIndexer::~DiskAnnIndexer() { destroy_io_ctx(init_ctx_); if (centroid_data_) { - free(centroid_data_); + DiskAnnUtil::free_aligned(centroid_data_); } - DiskAnnUtil::free_aligned(coord_cache_buf_); + reset_cache_storage(); } int DiskAnnIndexer::init(DiskAnnSearcherEntity &entity) { - entity_ = &entity; - auto storage = entity.get_storage(); auto vector_segment = entity.get_vector_segment(); + if (!storage || !vector_segment) { + LOG_ERROR("DiskAnn storage or vector segment is missing"); + return IndexError_InvalidFormat; + } + + const uint64_t stored_max_node_size = entity.max_node_size(); + if (stored_max_node_size == 0 || + stored_max_node_size > (std::numeric_limits::max)() || + stored_max_node_size < sizeof(uint32_t) || + meta_.element_size() > stored_max_node_size - sizeof(uint32_t)) { + LOG_ERROR("Invalid DiskAnn node size: node=%llu vector=%u", + static_cast(stored_max_node_size), + static_cast(meta_.element_size())); + return IndexError_InvalidFormat; + } + + const uint64_t stored_max_degree = entity.max_degree(); + const uint64_t neighbor_bytes = + stored_max_node_size - sizeof(uint32_t) - meta_.element_size(); + if (stored_max_degree > (std::numeric_limits::max)() || + stored_max_degree > neighbor_bytes / sizeof(diskann_id_t)) { + LOG_ERROR("Invalid DiskAnn node capacity: node=%llu vector=%u degree=%llu", + static_cast(stored_max_node_size), + static_cast(meta_.element_size()), + static_cast(stored_max_degree)); + return IndexError_InvalidFormat; + } + + const uint64_t expected_node_per_sector = + stored_max_node_size <= DiskAnnUtil::kSectorSize + ? DiskAnnUtil::kSectorSize / stored_max_node_size + : 0; + if (entity.node_per_sector() != expected_node_per_sector) { + LOG_ERROR( + "Invalid DiskAnn node layout: node=%llu nodes_per_sector=%llu " + "expected=%llu", + static_cast(stored_max_node_size), + static_cast(entity.node_per_sector()), + static_cast(expected_node_per_sector)); + return IndexError_InvalidFormat; + } + + uint64_t expected_index_size = 0; + if (entity.node_per_sector() > 0) { + const uint64_t sector_count = + entity.doc_cnt() / entity.node_per_sector() + + (entity.doc_cnt() % entity.node_per_sector() != 0 ? 1 : 0); + if (!checked_multiply_u64(sector_count, DiskAnnUtil::kSectorSize, + &expected_index_size)) { + LOG_ERROR("DiskAnn graph size overflows uint64"); + return IndexError_InvalidFormat; + } + } else { + const uint64_t sectors_per_node = + stored_max_node_size / DiskAnnUtil::kSectorSize + + (stored_max_node_size % DiskAnnUtil::kSectorSize != 0 ? 1 : 0); + uint64_t total_sectors = 0; + if (!checked_multiply_u64(entity.doc_cnt(), sectors_per_node, + &total_sectors) || + !checked_multiply_u64(total_sectors, DiskAnnUtil::kSectorSize, + &expected_index_size)) { + LOG_ERROR("DiskAnn graph size overflows uint64"); + return IndexError_InvalidFormat; + } + } + + const uint64_t vector_segment_size = vector_segment->data_size(); + const uint64_t vector_segment_offset = vector_segment->data_offset(); + if (entity.index_size() != expected_index_size || + vector_segment_size != expected_index_size || + vector_segment_offset % DiskAnnUtil::kSectorSize != 0) { + LOG_ERROR( + "Invalid DiskAnn graph layout: declared=%llu segment=%llu " + "expected=%llu offset=%llu", + static_cast(entity.index_size()), + static_cast(vector_segment_size), + static_cast(expected_index_size), + static_cast(vector_segment_offset)); + return IndexError_InvalidFormat; + } + + auto cached_file = storage->file(); + // DiskAnn must capture the exact file object that supplied every in-memory + // segment before releasing IndexStorage. FileReadStorage's + // alone_file_handle mode gives each Segment an independent file object and + // exposes no shared descriptor, so an atomic path replacement could mix + // metadata and graph data from different index snapshots. + if (!cached_file) { + LOG_ERROR( + "DiskAnn requires FileReadStorage with " + "proxima.file.read_storage.alone_file_handle disabled"); + return IndexError_InvalidArgument; + } + uint64_t graph_end = 0; + if (!checked_add_u64(vector_segment_offset, expected_index_size, + &graph_end) || + graph_end > cached_file->size()) { + LOG_ERROR( + "DiskAnn graph exceeds the captured file: end=%llu file_size=%llu", + static_cast(graph_end), + static_cast(cached_file->size())); + return IndexError_InvalidFormat; + } + + max_node_size_ = static_cast(stored_max_node_size); + sector_num_per_node_ = + DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); + if (sector_num_per_node_ == 0 || + beam_width_ > DiskAnnUtil::kMaxSectorReadNum / sector_num_per_node_) { + LOG_ERROR("DiskAnn node size exceeds the search buffer capacity"); + return IndexError_InvalidArgument; + } pq_table_ = entity.get_pq_table(); + entity_ = entity.clone(); + if (!entity_) { + LOG_ERROR("Failed to clone in-memory DiskAnn entity"); + return IndexError_NoMemory; + } index_segment_offset_ = vector_segment->data_offset(); - reader_.reset(new LinuxAlignedFileReader()); + const auto file_path = storage->file_path(); + int ret = 0; + reader_.reset(new PlatformAlignedFileReader()); +#if defined(_WIN32) || defined(_WIN64) + // Drop every Segment reference created by entity.load() before checking the + // File control block. Without an external alias, only cached_file and the + // FileReadStorage itself remain as owners. + entity.release_storage(); + vector_segment.reset(); + if (cached_file.use_count() != 2) { + LOG_ERROR( + "DiskAnn on Windows cannot load while the caller retains the " + "FileReadStorage file or one of its segments"); + return IndexError_InvalidArgument; + } + + // Capture the exact file object that supplied the in-memory metadata before + // releasing FileReadStorage. Reopening file_path after cleanup could bind + // graph reads to a replacement file while PQ/keys still belong to the old + // one. + ret = static_cast(reader_.get()) + ->open_from_handle(file_path, cached_file->native_handle()); +#else + // POSIX atomic replacement leaves an open descriptor bound to the old + // inode. Capture an independent descriptor before cleanup so graph reads + // use the same file object that supplied the in-memory metadata. + ret = static_cast(reader_.get()) + ->open_from_handle(file_path, cached_file->native_handle()); +#endif + if (ret != 0) { + LOG_ERROR("Failed to capture DiskAnn index file, ret=%d", ret); + return ret; + } - auto file_path = storage->file_path(); - reader_->open(file_path); + ret = storage->cleanup(); +#if !defined(_WIN32) && !defined(_WIN64) + entity.release_storage(); + vector_segment.reset(); +#endif + storage.reset(); + if (ret != 0) { + reader_->close(); + LOG_ERROR("Failed to release DiskAnn index storage, ret=%d", ret); + return ret; + } - storage->cleanup(); +#if defined(_WIN32) || defined(_WIN64) + // Windows cannot keep an ordinary buffered alias to this file object beside + // DiskAnn's unbuffered handles without a severe random-read regression. The + // preflight check above avoids consuming the storage on an ordinary + // ownership error. Check again after cleanup so an unexpected remaining + // owner cannot make the successful load retain a buffered handle. + if (cached_file.use_count() != 1) { + reader_->close(); + LOG_ERROR( + "DiskAnn on Windows cannot load while the caller retains the " + "FileReadStorage file or one of its segments"); + return IndexError_InvalidArgument; + } +#endif + // Releasing the last internal reference closes the buffered source handle. + // POSIX caller-owned aliases remain valid; Windows has rejected them above. + cached_file.reset(); - int ret = setup_io_ctx(init_ctx_); + ret = setup_io_ctx(init_ctx_); if (ret != 0) { LOG_ERROR("setup io ctx error"); return ret; } - max_node_size_ = entity.max_node_size(); disk_bytes_per_point_ = meta_.element_size(); node_per_sector_ = entity.node_per_sector(); - aligned_dim_ = meta_.dimension(); - pq_chunk_num_ = entity.pq_chunk_num(); medoid_ = entity.medoid(); entrypoints_.push_back(medoid_); - auto &entrypoints = entity.entrypoints(); + const auto &entrypoints = entity.entrypoints(); for (size_t i = 0; i < entrypoints.size(); ++i) { entrypoints_.push_back(entrypoints[i]); } doc_cnt_ = entity.doc_cnt(); - max_degree_ = entity.max_degree(); - - sector_num_per_node_ = - DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); - if (beam_width_ * sector_num_per_node_ > DiskAnnUtil::kMaxSectorReadNum) { - LOG_ERROR("Beamwidth can not be higher than kMaxSectorReadNum"); - - return IndexError_InvalidArgument; - } + max_degree_ = static_cast(stored_max_degree); centroid_stride_ = DiskAnnUtil::round_up(meta_.element_size(), 32); DiskAnnUtil::alloc_aligned(¢roid_data_, @@ -128,19 +340,89 @@ diskann_key_t DiskAnnIndexer::get_key(diskann_id_t id) const { return entity_->get_key(id); } +bool DiskAnnIndexer::should_include_result(DiskAnnContext *ctx, diskann_id_t id, + diskann_key_t *key) const { + const diskann_key_t resolved_key = get_key(id); + if (key != nullptr) { + *key = resolved_key; + } + return resolved_key != kInvalidKey && + (!ctx->filter().is_valid() || !ctx->filter()(resolved_key)); +} + diskann_id_t DiskAnnIndexer::get_id(diskann_key_t key) const { return entity_->get_id(key); } +int DiskAnnIndexer::parse_node_neighbors(const uint8_t *node_buf, + diskann_id_t node_id, + uint32_t &neighbor_count, + diskann_id_t *neighbors) const { + if (node_buf == nullptr) { + LOG_ERROR("DiskAnn node %u has no data buffer", node_id); + return IndexError_InvalidArgument; + } + + const uint8_t *neighbor_data = + DiskAnnUtil::offset_to_node_neighbor(node_buf, meta_.element_size()); + uint32_t parsed_count = 0; + memcpy(&parsed_count, neighbor_data, sizeof(parsed_count)); + if (parsed_count > max_degree_) { + LOG_ERROR("DiskAnn node %u has %u neighbors, exceeding max degree %u", + node_id, parsed_count, max_degree_); + return IndexError_InvalidFormat; + } + if (parsed_count != 0 && neighbors == nullptr) { + LOG_ERROR("DiskAnn node %u has no neighbor output buffer", node_id); + return IndexError_InvalidArgument; + } + + const uint8_t *neighbor_ids = neighbor_data + sizeof(parsed_count); + for (uint32_t i = 0; i < parsed_count; ++i) { + diskann_id_t neighbor_id = 0; + memcpy(&neighbor_id, neighbor_ids + i * sizeof(neighbor_id), + sizeof(neighbor_id)); + if (neighbor_id >= doc_cnt_) { + LOG_ERROR( + "DiskAnn node %u has invalid neighbor %u at position %u; " + "document count is %llu", + node_id, neighbor_id, i, static_cast(doc_cnt_)); + return IndexError_InvalidFormat; + } + neighbors[i] = neighbor_id; + } + + neighbor_count = parsed_count; + return 0; +} + std::vector DiskAnnIndexer::read_nodes( const std::vector &node_ids, std::vector &coord_buffers, std::vector> &neighbor_buffers) { - std::vector read_reqs; std::vector retval(node_ids.size(), true); + if (coord_buffers.size() != node_ids.size() || + neighbor_buffers.size() != node_ids.size()) { + LOG_ERROR( + "read_nodes: node, coordinate, and neighbor buffer counts must " + "match"); + std::fill(retval.begin(), retval.end(), false); + return retval; + } if (node_ids.empty()) { return retval; } + for (diskann_id_t node_id : node_ids) { + if (node_id >= doc_cnt_) { + LOG_ERROR("read_nodes: node %u exceeds document count %llu", node_id, + static_cast(doc_cnt_)); + std::fill(retval.begin(), retval.end(), false); + return retval; + } + } + + std::vector read_reqs; + read_reqs.reserve(node_ids.size()); uint8_t *buf = nullptr; auto sector_num = @@ -191,13 +473,15 @@ std::vector DiskAnnIndexer::read_nodes( } if (neighbor_buffers[i].second != nullptr) { - uint32_t *node_neighbor = - DiskAnnUtil::offset_to_node_neighbor(node_buf, meta_.element_size()); - uint32_t neighbor_num = *node_neighbor; + uint32_t neighbor_num = 0; + int parse_ret = parse_node_neighbors(node_buf, node_ids[i], neighbor_num, + neighbor_buffers[i].second); + if (parse_ret != 0) { + retval[i] = false; + continue; + } neighbor_buffers[i].first = neighbor_num; - memcpy(neighbor_buffers[i].second, node_neighbor + 1, - neighbor_num * sizeof(diskann_id_t)); } } @@ -206,77 +490,229 @@ std::vector DiskAnnIndexer::read_nodes( return retval; } -int DiskAnnIndexer::load_cache_list( - const std::vector &node_list) { - LOG_INFO("Loading the cache list into memory"); +void DiskAnnIndexer::reset_cache_storage() { + // The maps contain pointers into the two backing buffers. Drop the maps + // first so no stale pointer remains observable while storage is replaced. + coord_cache_.clear(); + neighbor_cache_.clear(); + DiskAnnUtil::free_aligned(coord_cache_buf_); + coord_cache_buf_ = nullptr; + std::vector().swap(neighbor_cache_buffer_); +} - size_t num_cached_nodes = node_list.size(); - if (num_cached_nodes == 0) { +uint32_t DiskAnnIndexer::effective_cache_node_count( + uint32_t requested_nodes) const { + uint64_t max_nodes = 0; + if (doc_cnt_ != 0) { + max_nodes = + doc_cnt_ / 10 + (doc_cnt_ % 10 >= 5 ? static_cast(1) : 0); + max_nodes = std::max(1, max_nodes); + } + const uint32_t effective_nodes = + static_cast(std::min(requested_nodes, max_nodes)); + if (effective_nodes != requested_nodes) { + LOG_WARN( + "Reducing nodes to cache from: %u, to: (10 percent of total nodes: " + "%u)", + requested_nodes, effective_nodes); + } + return effective_nodes; +} + +int DiskAnnIndexer::prepare_cache_storage(size_t capacity, + CacheLoadState &state) { + reset_cache_storage(); + state = {}; + state.capacity = capacity; + + if (capacity == 0) { return 0; } - neighbor_cache_buffer_.resize(num_cached_nodes * (max_degree_ + 1), 0); + const uint64_t neighbor_entries_per_node_u64 = + static_cast(max_degree_) + 1; + if (neighbor_entries_per_node_u64 > std::numeric_limits::max()) { + LOG_ERROR("DiskANN node cache neighbor stride overflow"); + return IndexError_InvalidArgument; + } + const size_t neighbor_entries_per_node = + static_cast(neighbor_entries_per_node_u64); + const size_t max_neighbor_entries = + std::numeric_limits::max() / sizeof(diskann_id_t); + if (capacity > max_neighbor_entries / neighbor_entries_per_node) { + LOG_ERROR("DiskANN node cache neighbor allocation size overflow"); + return IndexError_InvalidArgument; + } + + const size_t element_size = meta_.element_size(); + if (element_size == 0 || meta_.unit_size() == 0 || + capacity > std::numeric_limits::max() / element_size) { + LOG_ERROR("DiskANN node cache coordinate byte size overflow"); + reset_cache_storage(); + return IndexError_InvalidArgument; + } + const size_t coord_cache_bytes = capacity * element_size; + + try { + state.slots.reserve(capacity); + neighbor_cache_buffer_.resize(capacity * neighbor_entries_per_node, 0); + } catch (const std::exception &e) { + LOG_ERROR("Failed to allocate DiskANN node cache storage: %s", e.what()); + reset_cache_storage(); + state = {}; + return IndexError_NoMemory; + } - size_t coord_cache_buf_len = num_cached_nodes * aligned_dim_; - DiskAnnUtil::alloc_aligned((void **)&coord_cache_buf_, - coord_cache_buf_len * meta_.unit_size(), + DiskAnnUtil::alloc_aligned(&coord_cache_buf_, coord_cache_bytes, 8 * meta_.unit_size()); if (coord_cache_buf_ == nullptr) { LOG_ERROR("Failed to allocate coordinate cache buffer"); - neighbor_cache_buffer_.clear(); + reset_cache_storage(); return IndexError_NoMemory; } + return 0; +} - memset(coord_cache_buf_, 0, coord_cache_buf_len * meta_.unit_size()); - - constexpr size_t BLOCK_SIZE = 8; - size_t num_blocks = DiskAnnUtil::div_round_up(num_cached_nodes, BLOCK_SIZE); - for (size_t block = 0; block < num_blocks; block++) { - size_t start_idx = block * BLOCK_SIZE; - size_t end_idx = std::min(num_cached_nodes, (block + 1) * BLOCK_SIZE); +int DiskAnnIndexer::load_cache_list(CacheLoadState &state) { + LOG_INFO("Loading the remaining cache nodes into memory"); - std::vector nodes_to_read; - std::vector coord_buffers; - std::vector> neighbor_buffers; - for (size_t node_idx = start_idx; node_idx < end_idx; node_idx++) { - nodes_to_read.push_back(node_list[node_idx]); + std::vector pending_slots; + pending_slots.reserve(state.slots.size()); + for (size_t i = 0; i < state.slots.size(); ++i) { + if (!state.slots[i].loaded) { + pending_slots.push_back(i); + } + } + std::sort(pending_slots.begin(), pending_slots.end(), + [&state](size_t lhs, size_t rhs) { + return state.slots[lhs].id < state.slots[rhs].id; + }); + + const size_t neighbor_entries_per_node = static_cast(max_degree_) + 1; + const size_t batch_size = static_cast( + DiskAnnUtil::cache_load_batch_size(sector_num_per_node_)); + const size_t num_blocks = + DiskAnnUtil::div_round_up(pending_slots.size(), batch_size); + + std::vector nodes_to_read; + std::vector coord_buffers; + std::vector> neighbor_buffers; + nodes_to_read.reserve(batch_size); + coord_buffers.reserve(batch_size); + neighbor_buffers.reserve(batch_size); + + for (size_t block = 0; block < num_blocks; ++block) { + const size_t start_idx = block * batch_size; + const size_t end_idx = + std::min(pending_slots.size(), (block + 1) * batch_size); + + nodes_to_read.clear(); + coord_buffers.clear(); + neighbor_buffers.clear(); + for (size_t i = start_idx; i < end_idx; ++i) { + const size_t slot_idx = pending_slots[i]; + nodes_to_read.push_back(state.slots[slot_idx].id); coord_buffers.push_back(reinterpret_cast(coord_cache_buf_) + - node_idx * meta_.element_size()); + slot_idx * meta_.element_size()); neighbor_buffers.emplace_back( - 0, neighbor_cache_buffer_.data() + node_idx * (max_degree_ + 1)); + 0, + neighbor_cache_buffer_.data() + slot_idx * neighbor_entries_per_node); } - auto read_status = + const auto read_status = read_nodes(nodes_to_read, coord_buffers, neighbor_buffers); - - for (size_t i = 0; i < read_status.size(); i++) { - if (read_status[i] == true) { - coord_cache_.insert(std::make_pair(nodes_to_read[i], coord_buffers[i])); - neighbor_cache_.insert( - std::make_pair(nodes_to_read[i], neighbor_buffers[i])); + for (size_t i = 0; i < read_status.size(); ++i) { + if (read_status[i]) { + const size_t slot_idx = pending_slots[start_idx + i]; + state.slots[slot_idx].loaded = true; + state.slots[slot_idx].neighbor_count = neighbor_buffers[i].first; } } } - LOG_INFO("Load Cache List Done"); + // Publish both maps together only after all optional I/O has completed. + // Their values point into fixed-capacity buffers that will not move. + std::vector loaded_slots; + loaded_slots.reserve(state.slots.size()); + for (size_t i = 0; i < state.slots.size(); ++i) { + if (state.slots[i].loaded) { + loaded_slots.push_back(i); + } + } + std::sort(loaded_slots.begin(), loaded_slots.end(), + [&state](size_t lhs, size_t rhs) { + return state.slots[lhs].id < state.slots[rhs].id; + }); + + try { + for (size_t slot_idx : loaded_slots) { + const CacheSlot &slot = state.slots[slot_idx]; + void *coord = reinterpret_cast(coord_cache_buf_) + + slot_idx * meta_.element_size(); + diskann_id_t *neighbors = + neighbor_cache_buffer_.data() + slot_idx * neighbor_entries_per_node; + coord_cache_.emplace_hint(coord_cache_.end(), slot.id, coord); + neighbor_cache_.emplace_hint( + neighbor_cache_.end(), slot.id, + std::make_pair(slot.neighbor_count, neighbors)); + } + } catch (const std::exception &e) { + LOG_ERROR("Failed to publish DiskANN node cache: %s", e.what()); + reset_cache_storage(); + return IndexError_NoMemory; + } + + const size_t failed_nodes = state.slots.size() - loaded_slots.size(); + if (failed_nodes != 0) { + LOG_WARN( + "DiskANN node cache preload completed with read failures: " + "selected_nodes=%zu loaded_nodes=%zu failed_nodes=%zu", + state.slots.size(), loaded_slots.size(), failed_nodes); + } return 0; } -void DiskAnnIndexer::cache_bfs_levels(uint64_t num_nodes_to_cache, - std::vector &node_list) { - std::set node_set; +int DiskAnnIndexer::configure_cache(uint32_t cache_node_num) { + cache_node_num = effective_cache_node_count(cache_node_num); + if (cache_node_num == 0) { + reset_cache_storage(); + return 0; + } - size_t tenp_cnt = static_cast(std::round(doc_cnt_ * 0.1)); - if (num_nodes_to_cache > tenp_cnt) { - LOG_WARN( - "Reducing nodes to cache from: %zu, to: (10 percent of total nodes: " - "%zu)", - (size_t)num_nodes_to_cache, (size_t)tenp_cnt); + CacheLoadState state; + int ret = prepare_cache_storage(cache_node_num, state); + if (ret != 0) { + return ret; + } - num_nodes_to_cache = tenp_cnt == 0 ? 1 : tenp_cnt; + ailego::ElapsedTime cache_timer; + LOG_INFO("Caching %u nodes around medoid(s)", cache_node_num); + ret = cache_bfs_levels(cache_node_num, state); + if (ret != 0) { + reset_cache_storage(); + return ret; + } + ret = load_cache_list(state); + if (ret != 0) { + return ret; } + const size_t selected_nodes = state.slots.size(); + const size_t loaded_nodes = coord_cache_.size(); + LOG_INFO( + "Load Cache List Done: requested_nodes=%u selected_nodes=%zu " + "loaded_nodes=%zu failed_nodes=%zu elapsed_ms=%llu", + cache_node_num, selected_nodes, loaded_nodes, + selected_nodes - loaded_nodes, + static_cast(cache_timer.milli_seconds())); + return 0; +} + +int DiskAnnIndexer::cache_bfs_levels(uint64_t num_nodes_to_cache, + CacheLoadState &state) { + std::set node_set; + LOG_INFO("Begin to cache %zu Nodes", (size_t)num_nodes_to_cache); std::unordered_set cur_level; @@ -310,6 +746,15 @@ void DiskAnnIndexer::cache_bfs_levels(uint64_t num_nodes_to_cache, std::sort(nodes_to_expand.begin(), nodes_to_expand.end()); + if (nodes_to_expand.size() > state.capacity - state.slots.size()) { + LOG_ERROR("DiskANN node cache BFS exceeded its allocated capacity"); + return IndexError_Runtime; + } + const size_t first_slot = state.slots.size(); + for (diskann_id_t id : nodes_to_expand) { + state.slots.push_back(CacheSlot{id, 0, false}); + } + bool finish_flag = false; constexpr uint64_t BLOCK_SIZE = 1024; @@ -323,42 +768,42 @@ void DiskAnnIndexer::cache_bfs_levels(uint64_t num_nodes_to_cache, std::vector nodes_to_read(nodes_to_expand.begin() + start, nodes_to_expand.begin() + end); - std::vector coord_buffers(block_size, nullptr); + std::vector coord_buffers; + coord_buffers.reserve(block_size); - std::vector>> - neighbor_buffers; + std::vector> neighbor_buffers; neighbor_buffers.reserve(block_size); - + const size_t neighbor_entries_per_node = + static_cast(max_degree_) + 1; for (size_t i = 0; i < block_size; i++) { - neighbor_buffers.emplace_back( - 0, std::vector(max_degree_ + 1)); + const size_t slot_idx = first_slot + start + i; + coord_buffers.push_back(reinterpret_cast(coord_cache_buf_) + + slot_idx * meta_.element_size()); + neighbor_buffers.emplace_back(0, + neighbor_cache_buffer_.data() + + slot_idx * neighbor_entries_per_node); } - std::vector> neighbor_buffers_ptr; - neighbor_buffers_ptr.reserve(block_size); - for (size_t i = 0; i < block_size; i++) { - neighbor_buffers_ptr.emplace_back(neighbor_buffers[i].first, - neighbor_buffers[i].second.data()); - } + const auto read_status = + read_nodes(nodes_to_read, coord_buffers, neighbor_buffers); - auto read_status = - read_nodes(nodes_to_read, coord_buffers, neighbor_buffers_ptr); - - for (uint32_t i = 0; i < read_status.size(); i++) { - if (read_status[i] == false) { + for (size_t i = 0; i < read_status.size(); i++) { + if (!read_status[i]) { continue; - } else { - neighbor_buffers[i].first = neighbor_buffers_ptr[i].first; - uint32_t neighbor_num = neighbor_buffers[i].first; - diskann_id_t *neighbors = neighbor_buffers[i].second.data(); + } - for (uint32_t j = 0; j < neighbor_num && !finish_flag; j++) { - if (node_set.find(neighbors[j]) == node_set.end()) { - cur_level.insert(neighbors[j]); - } - if (cur_level.size() + node_set.size() >= num_nodes_to_cache) { - finish_flag = true; - } + const size_t slot_idx = first_slot + start + i; + state.slots[slot_idx].loaded = true; + state.slots[slot_idx].neighbor_count = neighbor_buffers[i].first; + + const uint32_t neighbor_num = neighbor_buffers[i].first; + diskann_id_t *neighbors = neighbor_buffers[i].second; + for (uint32_t j = 0; j < neighbor_num && !finish_flag; j++) { + if (node_set.find(neighbors[j]) == node_set.end()) { + cur_level.insert(neighbors[j]); + } + if (cur_level.size() + node_set.size() >= num_nodes_to_cache) { + finish_flag = true; } } } @@ -377,23 +822,22 @@ void DiskAnnIndexer::cache_bfs_levels(uint64_t num_nodes_to_cache, ailego_assert(node_set.size() + cur_level.size() == num_nodes_to_cache || cur_level.size() == 0); - node_list.clear(); - node_list.reserve(node_set.size() + cur_level.size()); - - for (auto node : node_set) { - node_list.push_back(node); + std::vector final_level(cur_level.begin(), cur_level.end()); + std::sort(final_level.begin(), final_level.end()); + if (final_level.size() > state.capacity - state.slots.size()) { + LOG_ERROR("DiskANN node cache frontier exceeded its allocated capacity"); + return IndexError_Runtime; } - - for (auto node : cur_level) { - node_list.push_back(node); + for (diskann_id_t id : final_level) { + state.slots.push_back(CacheSlot{id, 0, false}); } - size_t total_size = node_list.size(); + const size_t total_size = state.slots.size(); LOG_INFO("Level: %zu, Cached Size: %zu, Total Cached Size: %zu", (size_t)level, (size_t)(total_size - prev_node_set_size), (size_t)total_size); - return; + return 0; } int DiskAnnIndexer::linear_search(DiskAnnContext *ctx) { @@ -405,9 +849,13 @@ int DiskAnnIndexer::linear_search(DiskAnnContext *ctx) { auto &group_topk_heaps = ctx->group_topk_heaps(); group_topk_heaps.clear(); auto emplace_candidate = [&](diskann_id_t id, VectorInfo info) { + const diskann_key_t key = get_key(id); + if (key == kInvalidKey) { + return; + } if (ctx->group_by_search() && ctx->group_by().is_valid()) { topk_heap.emplace(id, info); - std::string group_id = ctx->group_by()(get_key(id)); + std::string group_id = ctx->group_by()(key); auto &group_topk_heap = group_topk_heaps[group_id]; if (group_topk_heap.empty()) { group_topk_heap.limit(ctx->group_topk()); @@ -452,7 +900,7 @@ int DiskAnnIndexer::linear_search(DiskAnnContext *ctx) { diskann_id_t id = 0; while (id < doc_cnt_) { while (frontier.size() < beam_width_) { - if (!ctx->filter().is_valid() || !ctx->filter()(get_key(id))) { + if (should_include_result(ctx, id, nullptr)) { auto iter = neighbor_cache_.find(id); if (iter != neighbor_cache_.end()) { cached_neighbors.push_back( @@ -555,9 +1003,13 @@ int DiskAnnIndexer::keys_search(const std::vector &keys, auto &group_topk_heaps = ctx->group_topk_heaps(); group_topk_heaps.clear(); auto emplace_candidate = [&](diskann_id_t id, VectorInfo info) { + const diskann_key_t key = get_key(id); + if (key == kInvalidKey) { + return; + } if (ctx->group_by_search() && ctx->group_by().is_valid()) { topk_heap.emplace(id, info); - std::string group_id = ctx->group_by()(get_key(id)); + std::string group_id = ctx->group_by()(key); auto &group_topk_heap = group_topk_heaps[group_id]; if (group_topk_heap.empty()) { group_topk_heap.limit(ctx->group_topk()); @@ -707,10 +1159,20 @@ int DiskAnnIndexer::keys_search(const std::vector &keys, int DiskAnnIndexer::get_vector(diskann_id_t id, IndexContext::Pointer &context, std::string &vector) { DiskAnnContext *ctx = dynamic_cast(context.get()); + if (ctx == nullptr) { + LOG_ERROR("get_vector: invalid DiskAnn context"); + return IndexError_InvalidArgument; + } auto &stats = ctx->query_stats(); IOContext &io_ctx = ctx->io_ctx(); + // Search contexts are owned by an external pool and may outlive this index. + // Fetch contexts, however, are owned by the provider/iterator that also owns + // the reader; retaining their private handle avoids reopening it per vector. + IOContextReleaseGuard release_guard( + *reader_, io_ctx, + ctx->context_type() == DiskAnnContext::kSearcherContext); uint8_t *sector_buffer = reinterpret_cast(ctx->sector_buffer()); @@ -718,6 +1180,23 @@ int DiskAnnIndexer::get_vector(diskann_id_t id, IndexContext::Pointer &context, node_per_sector_ > 0 ? 1 : DiskAnnUtil::div_round_up(max_node_size_, DiskAnnUtil::kSectorSize); + const size_t sector_read_size = + static_cast(sector_num_per_node) * DiskAnnUtil::kSectorSize; + const size_t node_offset = + node_per_sector_ == 0 + ? 0 + : static_cast(id % node_per_sector_) * max_node_size_; + if (sector_num_per_node == 0 || sector_buffer == nullptr || + sector_read_size > ctx->sector_buffer_size() || + node_offset > sector_read_size || + meta_.element_size() > sector_read_size - node_offset) { + LOG_ERROR( + "get_vector: invalid sector buffer range, read=%zu offset=%zu " + "vector=%u available=%zu", + sector_read_size, node_offset, + static_cast(meta_.element_size()), ctx->sector_buffer_size()); + return IndexError_InvalidArgument; + } ailego::ElapsedTime query_timer; ailego::ElapsedTime io_timer; @@ -755,8 +1234,7 @@ int DiskAnnIndexer::get_vector(diskann_id_t id, IndexContext::Pointer &context, DiskAnnUtil::get_node_sector(node_per_sector_, max_node_size_, DiskAnnUtil::kSectorSize, id) * DiskAnnUtil::kSectorSize, - sector_num_per_node * DiskAnnUtil::kSectorSize, - frontier_neighbor.second); + sector_read_size, frontier_neighbor.second); stats.disk_page_reads++; stats.io_num++; @@ -771,8 +1249,7 @@ int DiskAnnIndexer::get_vector(diskann_id_t id, IndexContext::Pointer &context, return IndexError_Runtime; } - uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( - node_per_sector_, max_node_size_, frontier_neighbor.second, id); + uint8_t *node_disk_buf = frontier_neighbor.second + node_offset; void *node_fp_coords = node_disk_buf; @@ -801,7 +1278,38 @@ int DiskAnnIndexer::knn_search(DiskAnnContext *ctx) { return 0; } +void DiskAnnIndexer::release_io_ctx(DiskAnnContext *ctx) { + if (reader_ && ctx) { + reader_->release_io_ctx(ctx->io_ctx()); + } +} + int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { + int error_code = IndexError_Runtime; + try { + return cached_beam_search_impl(ctx); + } catch (const std::bad_alloc &) { + LOG_ERROR("cached_beam_search: memory allocation failed"); + error_code = IndexError_NoMemory; + } catch (const std::exception &e) { + LOG_ERROR("cached_beam_search: unexpected exception: %s", e.what()); + } catch (...) { + LOG_ERROR("cached_beam_search: unknown exception"); + } + + // An exception may occur after an asynchronous batch has been submitted. + // Recreate the context only after destroy_io_ctx has cancelled and waited + // for all requests, so the caller can safely reuse or destroy this context. + IOContext &io_ctx = ctx->io_ctx(); + destroy_io_ctx(io_ctx); + if (setup_io_ctx(io_ctx) != 0) { + LOG_ERROR("cached_beam_search: failed to recreate I/O context"); + } + ctx->set_error(true); + return error_code; +} + +int DiskAnnIndexer::cached_beam_search_impl(DiskAnnContext *ctx) { auto &stats = ctx->query_stats(); auto &dc = ctx->dist_calculator(); auto &topk_heap = ctx->topk_heap(); @@ -873,6 +1381,7 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { cached_neighbors; cached_neighbors.reserve(2 * effective_beam_width); + std::vector parsed_neighbors(max_degree_); PendingBatch pending; while (candidates.has_unexpanded_node() && num_ios < io_limit_) { @@ -944,8 +1453,7 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { float cur_expanded_dist = dc.dist(ctx->query(), node_fp_coords_copy); - if (!ctx->filter().is_valid() || - !ctx->filter()(get_key(std::get<0>(cached_neighbor)))) { + if (should_include_result(ctx, std::get<0>(cached_neighbor), nullptr)) { topk_heap.emplace(std::get<0>(cached_neighbor), VectorInfo(cur_expanded_dist, make_vector_copy(node_fp_coords_copy))); @@ -976,6 +1484,7 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { if (!frontier.empty()) { std::vector completed; + int batch_parse_error = 0; while (pending.n_reaped < pending.n_submitted) { completed.clear(); io_timer.reset(); @@ -988,32 +1497,38 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { } for (uint32_t idx : completed) { + if (batch_parse_error != 0) { + continue; + } + auto &frontier_neighbor = frontier_neighbors[idx]; uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( node_per_sector_, max_node_size_, frontier_neighbor.second, frontier_neighbor.first); - uint32_t *node_buf = DiskAnnUtil::offset_to_node_neighbor( - node_disk_buf, meta_.element_size()); - uint32_t neighbor_num = *node_buf; + uint32_t neighbor_num = 0; + int parse_ret = + parse_node_neighbors(node_disk_buf, frontier_neighbor.first, + neighbor_num, parsed_neighbors.data()); + if (parse_ret != 0) { + batch_parse_error = parse_ret; + ctx->set_error(true); + continue; + } void *node_fp_coords = node_disk_buf; float cur_expanded_dist = dc.dist(ctx->query(), node_fp_coords); - if (!ctx->filter().is_valid() || - !ctx->filter()(get_key(frontier_neighbor.first))) { + if (should_include_result(ctx, frontier_neighbor.first, nullptr)) { topk_heap.emplace(frontier_neighbor.first, VectorInfo(cur_expanded_dist, make_vector_copy(node_fp_coords))); } - diskann_id_t *node_neighbors = - reinterpret_cast(node_buf + 1); - cpu_timer.reset(); std::vector distances(neighbor_num); - pq_table_->compute_dists(neighbor_num, node_neighbors, pq_chunk_num_, - ctx->pq_table_dist_buffer(), + pq_table_->compute_dists(neighbor_num, parsed_neighbors.data(), + pq_chunk_num_, ctx->pq_table_dist_buffer(), ctx->pq_coord_buffer(), distances.data()); stats.dist_num += neighbor_num; @@ -1021,7 +1536,7 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { cpu_timer.reset(); for (uint64_t m = 0; m < neighbor_num; ++m) { - diskann_id_t id = node_neighbors[m]; + diskann_id_t id = parsed_neighbors[m]; if (!visit_filter.visited(id)) { visit_filter.set_visited(id); stats.dist_num++; @@ -1033,6 +1548,9 @@ int DiskAnnIndexer::cached_beam_search(DiskAnnContext *ctx) { stats.cpu_us += cpu_timer.micro_seconds(); } } + if (batch_parse_error != 0) { + return batch_parse_error; + } } } @@ -1056,7 +1574,11 @@ void DiskAnnIndexer::populate_group_topk_heaps(DiskAnnContext *ctx) { for (uint32_t i = 0; i < topk_heap.size(); ++i) { diskann_id_t id = topk_heap[i].first; const auto &info = topk_heap[i].second; - std::string group_id = ctx->group_by()(get_key(id)); + const diskann_key_t key = get_key(id); + if (key == kInvalidKey) { + continue; + } + std::string group_id = ctx->group_by()(key); auto &group_topk_heap = group_topk_heaps[group_id]; if (group_topk_heap.empty()) { @@ -1109,9 +1631,6 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { : DiskAnnUtil::div_round_up( max_node_size_, DiskAnnUtil::kSectorSize); - pq_table_->preprocess_pq_dist_table(ctx->query_rotated(), - ctx->pq_table_dist_buffer()); - uint32_t num_ios = 0; std::vector frontier; @@ -1125,6 +1644,7 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { cached_neighbors.reserve(2 * beam_width_); uint64_t sector_buffer_idx; + std::vector parsed_neighbors(max_degree_); while (candidates.has_unexpanded_node() && num_ios < io_limit_) { frontier.clear(); @@ -1196,10 +1716,9 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { float cur_expanded_dist = dc.dist(ctx->query(), node_fp_coords_copy); - if (!ctx->filter().is_valid() || - !ctx->filter()(get_key(std::get<0>(cached_neighbor)))) { - std::string group_id = - ctx->group_by()(get_key(std::get<0>(cached_neighbor))); + diskann_key_t key = kInvalidKey; + if (should_include_result(ctx, std::get<0>(cached_neighbor), &key)) { + std::string group_id = ctx->group_by()(key); auto &group_topk_heap = group_topk_heaps[group_id]; if (group_topk_heap.empty()) { @@ -1242,19 +1761,23 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { uint8_t *node_disk_buf = DiskAnnUtil::offset_to_node( node_per_sector_, max_node_size_, frontier_neighbor.second, frontier_neighbor.first); - uint32_t *node_buf = DiskAnnUtil::offset_to_node_neighbor( - node_disk_buf, meta_.element_size()); - uint32_t neighbor_num = *node_buf; + uint32_t neighbor_num = 0; + int parse_ret = + parse_node_neighbors(node_disk_buf, frontier_neighbor.first, + neighbor_num, parsed_neighbors.data()); + if (parse_ret != 0) { + ctx->set_error(true); + return parse_ret; + } void *node_fp_coords = node_disk_buf; memcpy(data_buf, node_fp_coords, disk_bytes_per_point_); float cur_expanded_dist = dc.dist(ctx->query(), data_buf); - if (!ctx->filter().is_valid() || - !ctx->filter()(get_key(frontier_neighbor.first))) { - std::string group_id = - ctx->group_by()(get_key(frontier_neighbor.first)); + diskann_key_t key = kInvalidKey; + if (should_include_result(ctx, frontier_neighbor.first, &key)) { + std::string group_id = ctx->group_by()(key); auto &group_topk_heap = group_topk_heaps[group_id]; if (group_topk_heap.empty()) { @@ -1273,10 +1796,8 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { cpu_timer.reset(); std::vector distances(neighbor_num); - diskann_id_t *node_neighbors = - reinterpret_cast(node_buf + 1); - pq_table_->compute_dists(neighbor_num, node_neighbors, pq_chunk_num_, - ctx->pq_table_dist_buffer(), + pq_table_->compute_dists(neighbor_num, parsed_neighbors.data(), + pq_chunk_num_, ctx->pq_table_dist_buffer(), ctx->pq_coord_buffer(), distances.data()); stats.dist_num += neighbor_num; @@ -1284,7 +1805,7 @@ int DiskAnnIndexer::cached_beam_search_by_group(DiskAnnContext *ctx) { cpu_timer.reset(); for (uint64_t m = 0; m < neighbor_num; ++m) { - diskann_id_t id = node_neighbors[m]; + diskann_id_t id = parsed_neighbors[m]; visit_filter.set_visited(id); stats.dist_num++; diff --git a/src/core/algorithm/diskann/diskann_indexer.h b/src/core/algorithm/diskann/diskann_indexer.h index e3a27eb00..a3aaf3f9f 100644 --- a/src/core/algorithm/diskann/diskann_indexer.h +++ b/src/core/algorithm/diskann/diskann_indexer.h @@ -24,6 +24,9 @@ namespace zvec { namespace core { +class DiskAnnCacheTestPeer; +class DiskAnnStreamerTestPeer; + class DiskAnnIndexer { public: typedef std::shared_ptr Pointer; @@ -34,10 +37,8 @@ class DiskAnnIndexer { public: int init(DiskAnnSearcherEntity &entity); - int load_cache_list(const std::vector &node_list); - void cache_bfs_levels(uint64_t num_nodes_to_cache, - std::vector &node_list); + int configure_cache(uint32_t cache_node_num); int cached_beam_search(DiskAnnContext *ctx); int cached_beam_search_by_group(DiskAnnContext *ctx); @@ -48,6 +49,9 @@ class DiskAnnIndexer { int linear_search(DiskAnnContext *ctx); int keys_search(const std::vector &keys, DiskAnnContext *ctx); + //! Release lazy per-context reader resources at a public operation boundary. + void release_io_ctx(DiskAnnContext *ctx); + int get_vector(diskann_id_t id, IndexContext::Pointer &context, std::string &vector); @@ -69,9 +73,31 @@ class DiskAnnIndexer { void populate_group_topk_heaps(DiskAnnContext *ctx); private: - DiskAnnSearcherEntity *entity_; + struct CacheSlot { + diskann_id_t id{0}; + uint32_t neighbor_count{0}; + bool loaded{false}; + }; + + struct CacheLoadState { + size_t capacity{0}; + std::vector slots; + }; + + uint32_t effective_cache_node_count(uint32_t requested_nodes) const; + int prepare_cache_storage(size_t capacity, CacheLoadState &state); + int load_cache_list(CacheLoadState &state); + int cache_bfs_levels(uint64_t num_nodes_to_cache, CacheLoadState &state); + int parse_node_neighbors(const uint8_t *node_buf, diskann_id_t node_id, + uint32_t &neighbor_count, + diskann_id_t *neighbors) const; + bool should_include_result(DiskAnnContext *ctx, diskann_id_t id, + diskann_key_t *key) const; + void reset_cache_storage(); + int cached_beam_search_impl(DiskAnnContext *ctx); + + DiskAnnEntity::Pointer entity_{}; - IndexStorage::Pointer storage_{}; IndexMeta meta_; uint32_t max_degree_{0}; @@ -79,7 +105,6 @@ class DiskAnnIndexer { uint32_t max_node_size_{0}; uint64_t pq_chunk_num_{0}; uint64_t disk_bytes_per_point_{0}; - uint64_t aligned_dim_{0}; uint64_t index_segment_offset_{0}; uint64_t sector_num_per_node_{0}; @@ -89,22 +114,24 @@ class DiskAnnIndexer { diskann_id_t medoid_; std::vector entrypoints_; - std::shared_ptr reader_{nullptr}; + std::shared_ptr reader_{nullptr}; PQTable::Pointer pq_table_; - IOContext init_ctx_{0}; + IOContext init_ctx_{}; std::vector neighbor_cache_buffer_; void *coord_cache_buf_{nullptr}; std::map coord_cache_; std::map> neighbor_cache_; - uint32_t beam_width_{2}; uint32_t io_limit_{std::numeric_limits::max()}; uint64_t doc_cnt_{0}; + + friend class DiskAnnCacheTestPeer; + friend class DiskAnnStreamerTestPeer; }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_params.h b/src/core/algorithm/diskann/diskann_params.h index fac0dc60a..f4df80f5c 100644 --- a/src/core/algorithm/diskann/diskann_params.h +++ b/src/core/algorithm/diskann/diskann_params.h @@ -41,7 +41,6 @@ static const std::string PARAM_DISKANN_SEARCHER_LIST_SIZE( "zvec.diskann.searcher.list_size"); static const std::string PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM( "zvec.diskann.searcher.cache_node_num"); - static const std::string PARAM_DISKANN_REDUCER_INDEX_NAME( "zvec.diskann.reducer.index_name"); static const std::string PARAM_DISKANN_REDUCER_WORKING_PATH( diff --git a/src/core/algorithm/diskann/diskann_pq_table.cc b/src/core/algorithm/diskann/diskann_pq_table.cc index 0c13e4061..9e6edea80 100644 --- a/src/core/algorithm/diskann/diskann_pq_table.cc +++ b/src/core/algorithm/diskann/diskann_pq_table.cc @@ -24,9 +24,11 @@ PQTable::PQTable(const IndexMeta &meta, uint32_t chunk_num) if (meta.metric_name() == "Cosine") { if (meta.data_type() == IndexMeta::DataType::DT_FP32) { - meta_.set_dimension(meta.dimension() - 1); + meta_.set_dimension(meta.dimension() > 1 ? meta.dimension() - 1 : 0); + } else if (meta.data_type() == IndexMeta::DataType::DT_FP16) { + meta_.set_dimension(meta.dimension() > 2 ? meta.dimension() - 2 : 0); } else { - meta_.set_dimension(meta.dimension() - 2); + meta_.set_dimension(0); } } } @@ -37,13 +39,48 @@ int PQTable::init(std::vector &full_pivot_data, std::vector ¢roid, std::vector &chunk_offsets, std::vector &pq_data) { + if (meta_.data_type() != IndexMeta::DataType::DT_FP32 && + meta_.data_type() != IndexMeta::DataType::DT_FP16) { + LOG_ERROR("Unsupported DiskAnn PQ data type: %u", meta_.data_type()); + return IndexError_Unsupported; + } + if (meta_.dimension() == 0 || chunk_num_ == 0 || + chunk_num_ > meta_.dimension()) { + LOG_ERROR("Invalid DiskAnn PQ table dimensions"); + return IndexError_InvalidFormat; + } + + const size_t expected_pivot_size = + static_cast(meta_.element_size()) * kPQCentroidNum; + const size_t expected_centroid_size = meta_.element_size(); + if (full_pivot_data.size() != expected_pivot_size || + centroid.size() != expected_centroid_size || + chunk_offsets.size() != chunk_num_ + 1 || pq_data.empty() || + pq_data.size() % chunk_num_ != 0 || chunk_offsets.front() != 0 || + chunk_offsets.back() != meta_.dimension()) { + LOG_ERROR("Invalid DiskAnn PQ table buffer sizes or boundaries"); + return IndexError_InvalidFormat; + } + for (size_t i = 1; i < chunk_offsets.size(); ++i) { + if (chunk_offsets[i - 1] >= chunk_offsets[i] || + chunk_offsets[i] > meta_.dimension()) { + LOG_ERROR("Invalid DiskAnn PQ table chunk offsets"); + return IndexError_InvalidFormat; + } + } + full_pivot_data_ = std::move(full_pivot_data); centroid_ = std::move(centroid); chunk_offsets_ = std::move(chunk_offsets); pq_data_ = std::move(pq_data); // alloc and compute transpose - transposed_tables_.resize(kPQCentroidNum * meta_.element_size()); + try { + transposed_tables_.resize(expected_pivot_size); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate DiskAnn transposed PQ table"); + return IndexError_NoMemory; + } uint32_t dim = meta_.dimension(); uint32_t type = meta_.data_type(); diff --git a/src/core/algorithm/diskann/diskann_pq_trainer.cc b/src/core/algorithm/diskann/diskann_pq_trainer.cc index c84744cb8..787596475 100644 --- a/src/core/algorithm/diskann/diskann_pq_trainer.cc +++ b/src/core/algorithm/diskann/diskann_pq_trainer.cc @@ -13,14 +13,23 @@ // limitations under the License. #include "diskann_pq_trainer.h" +#include +#include +#include +#include +#include +#include +#include #include "diskann_entity.h" #include "diskann_util.h" namespace zvec { namespace core { -DiskAnnPqTrainer::DiskAnnPqTrainer(uint32_t max_train_sample_count) - : max_train_sample_count_{max_train_sample_count} {} +DiskAnnPqTrainer::DiskAnnPqTrainer(uint32_t max_train_sample_count, + double train_sample_ratio) + : max_train_sample_count_{max_train_sample_count}, + train_sample_ratio_{train_sample_ratio} {} DiskAnnPqTrainer::~DiskAnnPqTrainer() {} @@ -28,22 +37,46 @@ int DiskAnnPqTrainer::gen_random_sample(IndexHolder::Pointer holder, const IndexMeta &meta, std::string &sample_data, size_t &sample_size) { - double train_sample_ratio = - max_train_sample_count_ < 1 ? max_train_sample_count_ : 1; - - uint32_t max_train_sample_count = train_sample_ratio * holder->count(); - max_train_sample_count = max_train_sample_count > max_train_sample_count_ - ? max_train_sample_count_ - : max_train_sample_count; - - std::vector> sample_vecs; + sample_data.clear(); + sample_size = 0; + if (!holder || holder->count() == 0 || meta.element_size() == 0 || + max_train_sample_count_ == 0 || !std::isfinite(train_sample_ratio_) || + train_sample_ratio_ <= 0.0 || train_sample_ratio_ > 1.0) { + LOG_ERROR("Invalid DiskAnn PQ sampling configuration"); + return IndexError_InvalidArgument; + } - // Use a fixed seed for deterministic sampling across runs. - uint32_t x = 456321; - std::mt19937 gen(x); - std::uniform_real_distribution dist(0, 1); + const size_t holder_count = holder->count(); + long double requested_sample_count = + static_cast(holder_count) * train_sample_ratio_; + const long double nearest_integer = std::round(requested_sample_count); + const long double rounding_tolerance = + std::numeric_limits::epsilon() * + (std::max)(1.0L, std::fabs(requested_sample_count)); + if (std::fabs(requested_sample_count - nearest_integer) <= + rounding_tolerance) { + requested_sample_count = nearest_integer; + } + const size_t ratio_sample_count = static_cast( + std::ceil((std::min)(requested_sample_count, + static_cast(max_train_sample_count_)))); + const size_t target_sample_count = + (std::min)({holder_count, ratio_sample_count, + static_cast(max_train_sample_count_)}); + const size_t vec_size = meta.element_size(); + if (target_sample_count == 0 || + target_sample_count > (std::numeric_limits::max)() / vec_size) { + LOG_ERROR("Invalid DiskAnn PQ sample buffer size"); + return IndexError_InvalidLength; + } - uint32_t vec_size = meta.element_size(); + try { + sample_data.resize(target_sample_count * vec_size); + } catch (const std::bad_alloc &) { + return IndexError_NoMemory; + } catch (const std::length_error &) { + return IndexError_InvalidLength; + } auto iter = holder->create_iterator(); if (!iter) { @@ -51,33 +84,36 @@ int DiskAnnPqTrainer::gen_random_sample(IndexHolder::Pointer holder, return IndexError_Runtime; } - size_t sample_count = 0; - while (iter->is_valid() && sample_count < max_train_sample_count) { - float random = dist(gen); - - if (random < train_sample_ratio) { - const void *vec = iter->data(); - - std::vector temp_vec; - temp_vec.resize(vec_size); - - std::memcpy(reinterpret_cast(&temp_vec[0]), vec, vec_size); - - sample_vecs.push_back(std::move(temp_vec)); - - sample_count++; + // Reservoir sampling avoids the previous prefix bias while keeping builds + // deterministic for identical inputs. + std::mt19937_64 generator(456321); + size_t observed_count = 0; + while (iter->is_valid()) { + const void *vec = iter->data(); + if (vec == nullptr) { + LOG_ERROR("Failed to read a vector for DiskAnn PQ training"); + return IndexError_ReadData; } + size_t sample_index = observed_count; + if (observed_count >= target_sample_count) { + std::uniform_int_distribution distribution(0, observed_count); + sample_index = distribution(generator); + } + if (sample_index < target_sample_count) { + std::memcpy(sample_data.data() + sample_index * vec_size, vec, vec_size); + } + ++observed_count; iter->next(); } - sample_size = sample_vecs.size(); - sample_data.reserve(sample_size * vec_size); - - for (size_t i = 0; i < sample_size; i++) { - sample_data.append(reinterpret_cast(sample_vecs[i].data()), - vec_size); + sample_size = (std::min)(observed_count, target_sample_count); + if (sample_size == 0) { + sample_data.clear(); + LOG_ERROR("DiskAnn PQ training holder contains no vectors"); + return IndexError_InvalidLength; } + sample_data.resize(sample_size * vec_size); return 0; } @@ -89,11 +125,22 @@ int DiskAnnPqTrainer::prepare_pq_train_data( std::shared_ptr &train_features) { uint32_t dim = meta.dimension(); uint32_t vec_size = meta.element_size(); + if (num_train == 0 || dim == 0 || vec_size == 0 || + num_train > (std::numeric_limits::max)() / vec_size || + train_data.size() < num_train * vec_size || !train_features) { + LOG_ERROR("Invalid DiskAnn PQ training data"); + return IndexError_InvalidLength; + } - std::string train_data_processed; - train_data_processed.resize(num_train * vec_size); - - std::memcpy(&(train_data_processed[0]), train_data.data(), + std::vector train_data_processed; + try { + train_data_processed.resize(num_train * static_cast(dim)); + } catch (const std::bad_alloc &) { + return IndexError_NoMemory; + } catch (const std::length_error &) { + return IndexError_InvalidLength; + } + std::memcpy(train_data_processed.data(), train_data.data(), num_train * vec_size); // use fp32 to accumulate to avoid overflow @@ -102,7 +149,7 @@ int DiskAnnPqTrainer::prepare_pq_train_data( centroid_temp[d] = 0; } - T *train_data_processed_ptr = reinterpret_cast(&train_data_processed[0]); + T *train_data_processed_ptr = train_data_processed.data(); if (use_zero_mean) { for (uint64_t d = 0; d < dim; d++) { @@ -125,9 +172,9 @@ int DiskAnnPqTrainer::prepare_pq_train_data( // copy the centroid out centroid.resize(vec_size); - T *centroid_ptr = reinterpret_cast(centroid.data()); for (uint64_t d = 0; d < dim; d++) { - centroid_ptr[d] = centroid_temp[d]; + const T centroid_value = centroid_temp[d]; + std::memcpy(centroid.data() + d * sizeof(T), ¢roid_value, sizeof(T)); } return 0; @@ -149,19 +196,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, @@ -171,7 +225,7 @@ int DiskAnnPqTrainer::train_pq(IndexThreads::Pointer threads, std::vector ¢roid, std::vector &chunk_offsets) { uint32_t dim = meta.dimension(); - if (pq_chunk_num > dim) { + if (num_train == 0 || pq_chunk_num == 0 || pq_chunk_num > dim) { LOG_ERROR("Error: number of chunks more than dimension. chunk: %u, dim: %u", pq_chunk_num, dim); return IndexError_InvalidArgument; @@ -201,6 +255,10 @@ int DiskAnnPqTrainer::train_pq(IndexThreads::Pointer threads, return ret; } break; + + default: + LOG_ERROR("Unsupported DiskAnn PQ training data type: %u", type); + return IndexError_InvalidArgument; } // Do Train @@ -253,6 +311,9 @@ int DiskAnnPqTrainer::train_pq(IndexThreads::Pointer threads, return ret; } break; + + default: + return IndexError_InvalidArgument; } return 0; @@ -295,28 +356,51 @@ int DiskAnnPqTrainer::generate_pq(IndexThreads::Pointer threads, uint32_t pq_chunk_num, std::vector ¢roid, std::vector &block_compressed_data) { + if (!holder) { + LOG_ERROR("Cannot generate DiskAnn PQ data from a null holder"); + return IndexError_InvalidArgument; + } + uint32_t type = meta.data_type(); uint32_t dim = meta.dimension(); - if (pq_chunk_num > dim) { - LOG_ERROR("Error: number of chunks more than dimension. chunk: %u, dim: %u", - pq_chunk_num, dim); + const size_t element_size = meta.element_size(); + if (pq_chunk_num == 0 || pq_chunk_num > dim || element_size == 0 || + centroid.size() < element_size) { + LOG_ERROR( + "Invalid DiskAnn PQ generation metadata: chunk=%u dim=%u " + "element_size=%zu centroid_size=%zu", + pq_chunk_num, dim, element_size, centroid.size()); return IndexError_InvalidArgument; } // Do Label std::vector labels; size_t num_vecs = holder->count(); + if (num_vecs == 0) { + LOG_ERROR("Cannot generate DiskAnn PQ data for an empty holder"); + return IndexError_InvalidLength; + } size_t batch_size = num_vecs <= compress_batch_size_ ? num_vecs : compress_batch_size_; + if (num_vecs > (std::numeric_limits::max)() / pq_chunk_num || + batch_size > (std::numeric_limits::max)() / element_size) { + LOG_ERROR("DiskAnn PQ output size overflows"); + return IndexError_InvalidLength; + } - std::vector block_compressed_base(batch_size * pq_chunk_num); - - std::memset(&block_compressed_base[0], 0, - batch_size * pq_chunk_num * sizeof(uint32_t)); - - std::vector block_data(batch_size * meta.element_size()); - std::vector block_data_converted(batch_size * meta.element_size()); + std::vector fp32_block_data; + std::vector fp16_block_data; + switch (type) { + case IndexMeta::DataType::DT_FP32: + fp32_block_data.resize(batch_size * dim); + break; + case IndexMeta::DataType::DT_FP16: + fp16_block_data.resize(batch_size * dim); + break; + default: + return IndexError_InvalidArgument; + } size_t block_num = DiskAnnUtil::div_round_up(num_vecs, batch_size); @@ -334,42 +418,54 @@ int DiskAnnPqTrainer::generate_pq(IndexThreads::Pointer threads, size_t cur_block_size = end_id - start_id; - for (size_t i = 0; i < cur_block_size && iter->is_valid(); i++) { + for (size_t i = 0; i < cur_block_size; i++) { + if (!iter->is_valid()) { + LOG_ERROR("DiskAnn holder ended before its declared vector count"); + return IndexError_InvalidLength; + } const void *vec = iter->data(); - std::memcpy( - reinterpret_cast(&block_data[0]) + i * meta.element_size(), - vec, meta.element_size()); + if (vec == nullptr) { + LOG_ERROR("Failed to read a vector while generating DiskAnn PQ data"); + return IndexError_ReadData; + } + void *destination = + type == IndexMeta::DataType::DT_FP32 + ? static_cast(fp32_block_data.data() + i * dim) + : static_cast(fp16_block_data.data() + i * dim); + std::memcpy(destination, vec, element_size); iter->next(); } - std::memcpy(block_data_converted.data(), block_data.data(), - cur_block_size * meta.element_size()); - LOG_INFO("Processing Docs, Range: [%zu, %zu)..", start_id, end_id); std::shared_ptr block_features( new CompactIndexFeatures(meta)); switch (type) { - case IndexMeta::DataType::DT_FP32: + case IndexMeta::DataType::DT_FP32: { + std::vector typed_centroid(dim); + std::memcpy(typed_centroid.data(), centroid.data(), element_size); DiskAnnUtil::convert_vector_to_residual( - reinterpret_cast(block_data_converted.data()), - cur_block_size, dim, centroid.data()); + fp32_block_data.data(), cur_block_size, dim, typed_centroid.data()); + for (size_t i = 0; i < cur_block_size; i++) { + block_features->emplace(fp32_block_data.data() + i * dim); + } break; - case IndexMeta::DataType::DT_FP16: + } + case IndexMeta::DataType::DT_FP16: { + std::vector typed_centroid(dim); + std::memcpy(typed_centroid.data(), centroid.data(), element_size); DiskAnnUtil::convert_vector_to_residual( - reinterpret_cast(block_data_converted.data()), - cur_block_size, dim, centroid.data()); + fp16_block_data.data(), cur_block_size, dim, typed_centroid.data()); + for (size_t i = 0; i < cur_block_size; i++) { + block_features->emplace(fp16_block_data.data() + i * dim); + } break; + } default: return IndexError_InvalidArgument; } - for (size_t i = 0; i < cur_block_size; i++) { - block_features->emplace(block_data_converted.data() + - i * meta.element_size()); - } - int ret = chunk_cluster_.mount(block_features); if (ret != 0) { LOG_ERROR("Cannot mount block features"); @@ -392,6 +488,10 @@ int DiskAnnPqTrainer::generate_pq(IndexThreads::Pointer threads, LOG_INFO("Generate PQ Data Done."); } + if (iter->is_valid()) { + LOG_ERROR("DiskAnn holder contains more vectors than its declared count"); + return IndexError_InvalidLength; + } return 0; } diff --git a/src/core/algorithm/diskann/diskann_pq_trainer.h b/src/core/algorithm/diskann/diskann_pq_trainer.h index 8a83c28ef..0240cdeac 100644 --- a/src/core/algorithm/diskann/diskann_pq_trainer.h +++ b/src/core/algorithm/diskann/diskann_pq_trainer.h @@ -26,7 +26,8 @@ class DiskAnnPqTrainer { typedef std::unique_ptr UPointer; public: - DiskAnnPqTrainer(uint32_t max_train_sample_count); + DiskAnnPqTrainer(uint32_t max_train_sample_count, + double train_sample_ratio = PQTable::kTrainSampleRatio); virtual ~DiskAnnPqTrainer(); public: @@ -81,7 +82,8 @@ class DiskAnnPqTrainer { MultiChunkCluster chunk_cluster_; IndexCluster::CentroidList cluster_centroids_; uint32_t max_train_sample_count_{PQTable::kMaxTrainSampleCount}; + double train_sample_ratio_{PQTable::kTrainSampleRatio}; }; } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_reducer.cc b/src/core/algorithm/diskann/diskann_reducer.cc deleted file mode 100644 index b0eabcc5e..000000000 --- a/src/core/algorithm/diskann/diskann_reducer.cc +++ /dev/null @@ -1,190 +0,0 @@ -// 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 "diskann_reducer.h" -#include -#include -#include -#include -#include -#include "diskann_params.h" - -namespace zvec { -namespace core { - -int DiskAnnReducer::init(const ailego::Params ¶ms) { - params.get(PARAM_DISKANN_REDUCER_WORKING_PATH, &working_path_); - if (working_path_.empty()) { - LOG_ERROR("Missing parameter. %s", - PARAM_DISKANN_REDUCER_WORKING_PATH.c_str()); - return IndexError_InvalidArgument; - } - - std::string index_name = - params.get_as_string(PARAM_DISKANN_REDUCER_INDEX_NAME); - if (index_name.empty()) { - index_name = std::to_string(std::clock()); - } - - reducer_file_path_ = ailego::StringHelper::Concat( - working_path_, "/", kReducerFileName, index_name); - - holder_file_path_ = ailego::StringHelper::Concat(working_path_, "/", - kHolderFileName, index_name); - - state_ = STATE_INITED; - return 0; -} - -int DiskAnnReducer::cleanup(void) { - return 0; -} - -//! Reduce operator with filter -int DiskAnnReducer::reduce(const IndexFilter &filter) { - if (entities_.empty() || state_ != STATE_FEED) { - LOG_ERROR("No container to reduce, feed first"); - return IndexError_NoReady; - } - - if (use_mem_holder_) { - mem_holder_ = std::make_shared(meta_); - for (auto entity : entities_) { - size_t doc_cnt = entity->doc_cnt(); - for (size_t id = 0; id < doc_cnt; ++id) { - diskann_key_t pkey = entity->get_key(id); - - if (filter.is_valid() && filter(pkey)) { - continue; - } - - const void *vec = entity->get_vector(id); - mem_holder_->emplace(pkey, vec); - } - } - } else { - disk_holder_ = - std::make_shared(meta_, holder_file_path_); - - int ret = disk_holder_->init(); - if (ret != 0) { - LOG_ERROR("DiskAnn Index Holder init failed"); - return ret; - } - - for (auto entity : entities_) { - size_t doc_cnt = entity->doc_cnt(); - for (size_t id = 0; id < doc_cnt; ++id) { - diskann_key_t pkey = entity->get_key(id); - - if (filter.is_valid() && filter(pkey)) { - continue; - } - - const void *vec = entity->get_vector(id); - disk_holder_->emplace(pkey, vec); - } - } - - disk_holder_->close(); - } - - builder_ = IndexFactory::CreateBuilder(kDiskAnnBuilderName); - if (!builder_) { - LOG_ERROR("Create builder failed. name[%s]", kDiskAnnBuilderName.c_str()); - return IndexError_Runtime; - } - - if (thread_pool_ == nullptr) { - LOG_ERROR( - "Only support multi-thread mode. Thread pool is not set for reducer."); - return IndexError_NoReady; - } - - LOG_INFO("Start diskann reduce"); - - ailego::ElapsedTime timer; - - auto params = meta_.builder_params(); - - int ret = builder_->init(meta_, params); - if (ret != 0) { - LOG_ERROR("Init proxima streamer failed. ret[%d]", ret); - return ret; - } - - if (use_mem_holder_) { - ret = builder_->train(mem_holder_); - if (ret != 0) { - LOG_ERROR("Diskann builder failed to train. ret[%d]", ret); - return ret; - } - - ret = builder_->build(mem_holder_); - if (ret != 0) { - LOG_ERROR("Diskann builder failed to build. ret[%d]", ret); - return ret; - } - } else { - ret = builder_->train(disk_holder_); - if (ret != 0) { - LOG_ERROR("Diskann builder failed to train. ret[%d]", ret); - return ret; - } - - ret = builder_->build(disk_holder_); - if (ret != 0) { - LOG_ERROR("Diskann builder failed to build. ret[%d]", ret); - return ret; - } - } - - auto &stats = builder_->stats(); - - stats_.set_reduced_costtime(timer.seconds()); - stats_.set_filtered_count(stats.discarded_count()); - - state_ = STATE_REDUCE; - - LOG_INFO("End DiskAnn reduce. cost time: [%zu]s", (size_t)timer.seconds()); - return 0; -} - -//! Dump index by dumper -int DiskAnnReducer::dump(const IndexDumper::Pointer &dumper) { - LOG_INFO("Begin diskann reducer dump"); - - if (state_ != STATE_REDUCE) { - LOG_WARN("Reduce first before dump."); - return IndexError_NoReady; - } - - ailego::ElapsedTime timer; - int ret = builder_->dump(dumper); - if (ret != 0) { - LOG_ERROR("diskann reducer dump failed. ret[%d]", ret); - return ret; - } - stats_.set_dumped_costtime(timer.seconds()); - - LOG_INFO("End diskann reducer dump, dump costtime=[%zu]s", - (size_t)(stats_.dumped_costtime())); - - return 0; -} - -INDEX_FACTORY_REGISTER_REDUCER(DiskAnnReducer); - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_reducer.h b/src/core/algorithm/diskann/diskann_reducer.h deleted file mode 100644 index 60e40297b..000000000 --- a/src/core/algorithm/diskann/diskann_reducer.h +++ /dev/null @@ -1,85 +0,0 @@ -// 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. -#pragma once - -#include -#include -#include -#include -#include -#include -#include "diskann_holder.h" -#include "diskann_reducer_entity.h" - -namespace zvec { -namespace core { - -class DiskAnnReducer : public IndexReducer { - public: - //! Constructor - DiskAnnReducer(void) = default; - - protected: - //! Initialize Reducer - int init(const ailego::Params ¶ms) override; - - //! Cleanup Reducer - int cleanup(void) override; - - //! Feed indexes from containers - // int feed(IndexStorage::Pointer container) override; - - //! Reduce operator (with filter) - int reduce(const IndexFilter &filter) override; - - //! Dump index by dumper - int dump(const IndexDumper::Pointer &dumper) override; - - //! Retrieve statistics - const Stats &stats(void) const override { - return stats_; - } - - private: - enum State { - STATE_UNINITED = 0, - STATE_INITED = 1, - STATE_FEED = 2, - STATE_REDUCE = 3 - }; - - std::string working_path_{""}; - - IndexMeta meta_{}; - std::vector entities_{}; - - // bool use_mem_holder_{true}; - bool use_mem_holder_{false}; - RandomAccessIndexHolder::Pointer mem_holder_; - DiskAnnIndexHolder::Pointer disk_holder_; - - IndexBuilder::Pointer builder_{nullptr}; - std::string reducer_file_path_{""}; - std::string holder_file_path_{""}; - - Stats stats_{}; - State state_{STATE_UNINITED}; - - const std::string kDiskAnnBuilderName{"DiskAnnBuilder"}; - const std::string kReducerFileName{"diskann.reducer.builder."}; - const std::string kHolderFileName{"diskann.reducer.holder."}; -}; - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_reducer_entity.cc b/src/core/algorithm/diskann/diskann_reducer_entity.cc deleted file mode 100644 index 4ccfb6d21..000000000 --- a/src/core/algorithm/diskann/diskann_reducer_entity.cc +++ /dev/null @@ -1,215 +0,0 @@ -// 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 "diskann_reducer_entity.h" -#include -#include - -namespace zvec { -namespace core { - -int DiskAnnReducerEntity::load(const IndexStorage::Pointer &container, - bool check_crc) { - container_ = container; - - int ret = load_segments(check_crc); - if (ret != 0) { - return ret; - } - - sector_num_per_node_ = node_per_sector() > 0 - ? 1 - : DiskAnnUtil::div_round_up( - max_node_size(), DiskAnnUtil::kSectorSize); - - loaded_ = true; - - return 0; -} - -int DiskAnnReducerEntity::load_segments(bool /*check_crc*/) { - int ret; - ret = load_header_segment(); - if (ret != 0) { - LOG_ERROR("Load Header Segment Failed, ret = %d", ret); - - return ret; - } - - ret = load_key_segment(); - if (ret != 0) { - LOG_ERROR("Load Key Segment Failed, ret = %d", ret); - - return ret; - } - - ret = load_vector_segment(); - if (ret != 0) { - LOG_ERROR("Load Vector Segment Failed, ret = %d", ret); - - return ret; - } - - return 0; -} - -int DiskAnnReducerEntity::load_header_segment() { - const void *data = nullptr; - meta_segment_ = container_->get(kDiskAnnMetaSegmentId); - if (!meta_segment_ || - meta_segment_->data_size() < sizeof(DiskAnnMetaHeader)) { - LOG_ERROR("Miss or invalid segment %s", kDiskAnnMetaSegmentId.c_str()); - return IndexError_InvalidFormat; - } - if (meta_segment_->read(0, reinterpret_cast(&data), - sizeof(DiskAnnMetaHeader)) != - sizeof(DiskAnnMetaHeader)) { - LOG_ERROR("Read segment %s failed", kDiskAnnMetaSegmentId.c_str()); - return IndexError_ReadData; - } - - ::memcpy(reinterpret_cast(&meta_header_), data, - sizeof(DiskAnnMetaHeader)); - - return 0; -} - -int DiskAnnReducerEntity::load_vector_segment() { - vector_segment_ = container_->get(kDiskAnnVectorSegmentId); - if (!vector_segment_) { - LOG_ERROR("Miss or invalid segment %s", - DiskAnnEntity::kDiskAnnVectorSegmentId.c_str()); - return IndexError_InvalidFormat; - } - - return 0; -} - -int DiskAnnReducerEntity::load_key_segment() { - // load key - key_segment_ = container_->get(kDiskAnnKeySegmentId); - if (!key_segment_) { - LOG_ERROR("Miss or invalid segment %s", - DiskAnnEntity::kDiskAnnKeySegmentId.c_str()); - return IndexError_InvalidFormat; - } - - size_t key_data_len = doc_cnt() * sizeof(key_t); - - // load key mapping - key_mapping_segment_ = container_->get(kDiskAnnKeyMappingSegmentId); - const void *data = nullptr; - if (key_mapping_segment_->read(0, reinterpret_cast(&data), - key_data_len) != key_data_len) { - LOG_ERROR("Read segment %s failed", kDiskAnnKeyMappingSegmentId.c_str()); - return IndexError_ReadData; - } - - key_buffer_.resize(key_data_len); - memcpy(&(key_buffer_[0]), data, key_data_len); - - return 0; -} - -bool DiskAnnReducerEntity::do_crc_check( - std::vector &segments) const { - constexpr size_t blk_size = 4096; - const void *data; - - for (auto &segment : segments) { - size_t offset = 0; - size_t rd_size; - uint32_t crc = 0; - while (offset < segment->data_size()) { - size_t size = std::min(blk_size, segment->data_size() - offset); - if ((rd_size = segment->read(offset, &data, size)) <= 0) { - break; - } - offset += rd_size; - crc = ailego::Crc32c::Hash(data, rd_size, crc); - } - if (crc != segment->data_crc()) { - return false; - } - } - return true; -} - -//! Get vector local id by key -diskann_id_t DiskAnnReducerEntity::get_id(diskann_key_t key) const { - const diskann_id_t *key_mapping_data_ptr = - reinterpret_cast(key_mapping_buffer_.data()); - const diskann_key_t *key_data_ptr = - reinterpret_cast(key_buffer_.data()); - - //! Do binary search - diskann_id_t start = 0UL; - diskann_id_t end = doc_cnt(); - diskann_id_t idx = 0u; - while (start < end) { - idx = start + (end - start) / 2; - diskann_id_t local_id = key_mapping_data_ptr[idx]; - - const diskann_key_t local_key = key_data_ptr[local_id]; - - if (local_key < key) { - start = idx + 1; - } else if (local_key > key) { - end = idx; - } else { - return local_id; - } - } - - return kInvalidId; -} - -diskann_key_t DiskAnnReducerEntity::get_key(diskann_id_t id) const { - const void *key; - if (ailego_unlikely(key_segment_->read(id * sizeof(diskann_key_t), &key, - sizeof(diskann_key_t)) != - sizeof(diskann_key_t))) { - LOG_ERROR("Read key from segment failed"); - return kInvalidKey; - } - - return *(reinterpret_cast(key)); -} - -const void *DiskAnnReducerEntity::get_vector(diskann_id_t id) const { - size_t read_size = sector_num_per_node_ * DiskAnnUtil::kSectorSize; - size_t sector_id = DiskAnnUtil::get_node_sector( - node_per_sector(), max_node_size(), DiskAnnUtil::kSectorSize, id); - size_t offset = sector_id * DiskAnnUtil::kSectorSize; - - if (sector_id != sector_id_) { - const void *sector_data; - if (ailego_unlikely(vector_segment_->read(offset, §or_data, - read_size) != read_size)) { - LOG_ERROR("Read vector from segment failed"); - return nullptr; - } - - sector_id_ = sector_id; - sector_buffer_.assign(reinterpret_cast(sector_data), - read_size); - } - - return DiskAnnUtil::offset_to_node_const( - node_per_sector(), max_node_size(), - reinterpret_cast(sector_buffer_.data()), id); -} - -} // namespace core -} // namespace zvec diff --git a/src/core/algorithm/diskann/diskann_reducer_entity.h b/src/core/algorithm/diskann/diskann_reducer_entity.h deleted file mode 100644 index 3ab893045..000000000 --- a/src/core/algorithm/diskann/diskann_reducer_entity.h +++ /dev/null @@ -1,69 +0,0 @@ -// 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. -#pragma once - -#include -#include -#include -#include -#include "diskann_entity.h" -#include "diskann_file_reader.h" -#include "diskann_pq_table.h" -#include "diskann_util.h" - -namespace zvec { -namespace core { - -class DiskAnnReducerEntity : public DiskAnnEntity { - public: - using Pointer = std::shared_ptr; - using SegmentPointer = IndexStorage::Segment::Pointer; - - public: - DiskAnnReducerEntity() = default; - virtual ~DiskAnnReducerEntity() = default; - - int load(const IndexStorage::Pointer &container, bool check_crc); - int load_segments(bool check_crc); - int load_header_segment(); - int load_vector_segment(); - int load_key_segment(); - int load_key_mapping_segment(); - - bool do_crc_check(std::vector &segments) const; - - diskann_id_t get_id(diskann_key_t key) const override; - diskann_key_t get_key(diskann_id_t id) const override; - const void *get_vector(diskann_id_t id) const override; - - private: - IndexStorage::Pointer container_{}; - IndexStorage::Segment::Pointer meta_segment_{}; - IndexStorage::Segment::Pointer vector_segment_{}; - IndexStorage::Segment::Pointer key_segment_{}; - IndexStorage::Segment::Pointer key_mapping_segment_{}; - - std::string key_buffer_; - std::string key_mapping_buffer_; - - size_t sector_num_per_node_{0}; - - mutable size_t sector_id_{-1U}; - mutable std::string sector_buffer_; - - bool loaded_{false}; -}; - -} // namespace core -} // namespace zvec \ No newline at end of file diff --git a/src/core/algorithm/diskann/diskann_searcher.cc b/src/core/algorithm/diskann/diskann_searcher.cc index 630f5d294..c08a4b7ff 100644 --- a/src/core/algorithm/diskann/diskann_searcher.cc +++ b/src/core/algorithm/diskann/diskann_searcher.cc @@ -13,6 +13,9 @@ // limitations under the License. #include "diskann_searcher.h" +#include +#include +#include #include "diskann_context.h" #include "diskann_indexer.h" #include "diskann_params.h" @@ -45,13 +48,32 @@ int DiskAnnSearcher::init(const ailego::Params &search_params) { return IndexError_NoReady; } - params_ = search_params; - list_size_ = 200; - cache_nodes_num_ = 0; log_diskann_io_backend(); - params_.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); - params_.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); + uint32_t list_size = 200; + uint32_t cache_nodes_num = 0; + search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size); + if (list_size == 0) { + LOG_ERROR("list_size must be positive"); + return IndexError_InvalidArgument; + } + long long configured_cache_nodes = 0; + if (search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, + &configured_cache_nodes)) { + if (configured_cache_nodes < 0 || + static_cast(configured_cache_nodes) > + std::numeric_limits::max()) { + LOG_ERROR("cache_node_num must be in [0, UINT32_MAX]"); + return IndexError_InvalidArgument; + } + cache_nodes_num = static_cast(configured_cache_nodes); + } + + // Commit only after every value has been validated. A failed re-init must + // leave either the previous valid configuration or STATE_INIT untouched. + params_ = search_params; + list_size_ = list_size; + cache_nodes_num_ = cache_nodes_num; state_ = STATE_INITED; return 0; } @@ -84,6 +106,12 @@ int DiskAnnSearcher::load(IndexStorage::Pointer storage, LOG_ERROR("Initialize and unload DiskAnnSearcher before loading an index"); return IndexError_NoReady; } + if (!storage->file()) { + LOG_ERROR( + "DiskAnn requires storage with a shared file handle; disable " + "proxima.file.read_storage.alone_file_handle"); + return IndexError_InvalidArgument; + } diskann_indexer_.reset(); entity_.clear(); @@ -112,23 +140,13 @@ int DiskAnnSearcher::load(IndexStorage::Pointer storage, return res; } - if (cache_nodes_num_ != 0) { - std::vector node_list; - LOG_INFO("Caching %u nodes around medoid(s)", cache_nodes_num_); - - diskann_indexer_->cache_bfs_levels(cache_nodes_num_, node_list); - - ret = diskann_indexer_->load_cache_list(node_list); - if (ret != 0) { - return ret; - } - - node_list.clear(); - node_list.shrink_to_fit(); + ret = diskann_indexer_->configure_cache(cache_nodes_num_); + if (ret != 0) { + return ret; } if (measure) { - measure_ = measure; + measure_ = std::move(measure); } else { measure_ = IndexFactory::CreateMetric(meta_.metric_name()); if (!measure_) { @@ -229,6 +247,7 @@ int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, if (ret != 0) { return ret; } + AILEGO_DEFER(diskann_indexer_.get(), &DiskAnnIndexer::release_io_ctx, ctx); if (ailego_unlikely(!group_options_valid(ctx))) { LOG_ERROR("Group search requires a callback and a positive group topk"); return IndexError_InvalidArgument; @@ -238,6 +257,7 @@ int DiskAnnSearcher::search_impl(const void *query, const IndexQueryMeta &qmeta, ctx->resize_results(count); for (uint32_t i = 0; i < count; i++) { + ctx->visit_filter().clear(); ctx->reset_query(query); ret = diskann_indexer_->knn_search(ctx); @@ -283,6 +303,7 @@ int DiskAnnSearcher::search_bf_impl(const void *query, if (ret != 0) { return ret; } + AILEGO_DEFER(diskann_indexer_.get(), &DiskAnnIndexer::release_io_ctx, ctx); if (ailego_unlikely(!group_options_valid(ctx))) { LOG_ERROR("Group search requires a callback and a positive group topk"); return IndexError_InvalidArgument; @@ -343,6 +364,7 @@ int DiskAnnSearcher::search_bf_by_p_keys_impl( if (ret != 0) { return ret; } + AILEGO_DEFER(diskann_indexer_.get(), &DiskAnnIndexer::release_io_ctx, ctx); if (ailego_unlikely(!group_options_valid(ctx))) { LOG_ERROR("Group search requires a callback and a positive group topk"); return IndexError_InvalidArgument; diff --git a/src/core/algorithm/diskann/diskann_searcher.h b/src/core/algorithm/diskann/diskann_searcher.h index ee9900d56..fb7e54b52 100644 --- a/src/core/algorithm/diskann/diskann_searcher.h +++ b/src/core/algorithm/diskann/diskann_searcher.h @@ -17,11 +17,11 @@ #include "diskann_context.h" #include "diskann_indexer.h" -class LinuxAlignedFileReader; - namespace zvec { namespace core { +class DiskAnnCacheTestPeer; + class DiskAnnSearcher : public IndexSearcher { public: using ContextPointer = IndexSearcher::Context::Pointer; @@ -41,7 +41,8 @@ class DiskAnnSearcher : public IndexSearcher { int cleanup(void) override; //! Load Index from storage - int load(IndexStorage::Pointer storage, IndexMetric::Pointer metric) override; + int load(IndexStorage::Pointer storage, + IndexMetric::Pointer /*metric*/) override; //! Unload index from storage int unload(void) override; @@ -152,8 +153,6 @@ class DiskAnnSearcher : public IndexSearcher { uint32_t list_size_{200}; uint32_t cache_nodes_num_{0}; - bool warm_up_{false}; - uint32_t beam_size_{2}; DiskAnnIndexer::Pointer diskann_indexer_{nullptr}; DiskAnnSearcherEntity entity_{}; @@ -162,6 +161,8 @@ class DiskAnnSearcher : public IndexSearcher { Stats stats_; State state_{STATE_INIT}; + + friend class DiskAnnCacheTestPeer; }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_searcher_entity.cc b/src/core/algorithm/diskann/diskann_searcher_entity.cc index 7f4074fef..28cb8c8c0 100644 --- a/src/core/algorithm/diskann/diskann_searcher_entity.cc +++ b/src/core/algorithm/diskann/diskann_searcher_entity.cc @@ -13,11 +13,82 @@ // limitations under the License. #include "diskann_searcher_entity.h" +#include +#include +#include namespace zvec { namespace core { +namespace { + +bool checked_multiply(size_t lhs, size_t rhs, size_t *result) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return false; + } + *result = lhs * rhs; + return true; +} + +bool checked_add(size_t lhs, size_t rhs, size_t *result) { + if (rhs > std::numeric_limits::max() - lhs) { + return false; + } + *result = lhs + rhs; + return true; +} + +int get_pq_layout(const IndexMeta &meta, uint32_t *dimension, + size_t *full_pivot_size, size_t *centroid_size) { + uint32_t pq_dimension = meta.dimension(); + size_t unit_size = 0; + switch (meta.data_type()) { + case IndexMeta::DataType::DT_FP32: + unit_size = sizeof(float); + if (meta.metric_name() == "Cosine") { + if (pq_dimension <= 1) { + return IndexError_InvalidFormat; + } + --pq_dimension; + } + break; + case IndexMeta::DataType::DT_FP16: + unit_size = sizeof(ailego::Float16); + if (meta.metric_name() == "Cosine") { + if (pq_dimension <= 2) { + return IndexError_InvalidFormat; + } + pq_dimension -= 2; + } + break; + default: + return IndexError_Unsupported; + } + + if (pq_dimension == 0 || + !checked_multiply(pq_dimension, unit_size, centroid_size) || + !checked_multiply(*centroid_size, PQTable::kPQCentroidNum, + full_pivot_size)) { + return IndexError_InvalidFormat; + } + *dimension = pq_dimension; + return 0; +} + +} // namespace + void DiskAnnSearcherEntity::clear() { + release_storage(); + pq_table_.reset(); + key_buffer_.reset(); + key_mapping_buffer_.reset(); + entrypoints_.reset(); + meta_.clear(); + meta_header_ = {}; + pq_meta_ = {}; +} + +void DiskAnnSearcherEntity::release_storage() { storage_.reset(); meta_segment_.reset(); pq_meta_segment_.reset(); @@ -26,68 +97,30 @@ void DiskAnnSearcherEntity::clear() { key_segment_.reset(); key_mapping_segment_.reset(); entrypoint_segment_.reset(); - pq_table_.reset(); - key_buffer_.clear(); - key_mapping_buffer_.clear(); - entrypoints_.clear(); - meta_.clear(); - meta_header_ = {}; - pq_meta_ = {}; } const DiskAnnEntity::Pointer DiskAnnSearcherEntity::clone() const { - auto meta_segment = meta_segment_->clone(); - if (ailego_unlikely(!meta_segment)) { - LOG_ERROR("clone segment %s failed", kDiskAnnMetaSegmentId.c_str()); - return DiskAnnEntity::Pointer(); - } - - auto pq_meta_segment = pq_meta_segment_->clone(); - if (ailego_unlikely(!pq_meta_segment)) { - LOG_ERROR("clone segment %s failed", kDiskAnnPqMetaSegmentId.c_str()); - return DiskAnnEntity::Pointer(); - } - - auto pq_data_segment = pq_data_segment_->clone(); - if (ailego_unlikely(!pq_data_segment)) { - LOG_ERROR("clone segment %s failed", kDiskAnnPqDataSegmentId.c_str()); - return DiskAnnEntity::Pointer(); - } - - auto vector_segment = vector_segment_->clone(); - if (ailego_unlikely(!vector_segment)) { - LOG_ERROR("clone segment %s failed", kDiskAnnVectorSegmentId.c_str()); - return DiskAnnEntity::Pointer(); - } - - auto key_segment = key_segment_->clone(); - if (ailego_unlikely(!key_segment)) { - LOG_ERROR("clone segment %s failed", kDiskAnnKeySegmentId.c_str()); - return DiskAnnEntity::Pointer(); - } - - auto key_mapping_segment = key_mapping_segment_->clone(); - if (ailego_unlikely(!key_mapping_segment)) { - LOG_ERROR("clone segment %s failed", kDiskAnnKeyMappingSegmentId.c_str()); + std::unique_ptr entity(new (std::nothrow) + DiskAnnSearcherEntity()); + if (ailego_unlikely(!entity)) { + LOG_ERROR("DiskAnnSearcherEntity new failed"); return DiskAnnEntity::Pointer(); } - auto entrypoint_segment = entrypoint_segment_->clone(); - if (ailego_unlikely(!entrypoint_segment)) { - LOG_ERROR("clone segment %s failed", kDiskAnnEntryPointSegmentId.c_str()); + try { + entity->meta_header_ = meta_header_; + entity->pq_meta_ = pq_meta_; + entity->meta_ = meta_; + entity->pq_table_ = pq_table_; + entity->key_buffer_ = key_buffer_; + entity->key_mapping_buffer_ = key_mapping_buffer_; + entity->entrypoints_ = entrypoints_; + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to clone in-memory DiskAnn entity"); return DiskAnnEntity::Pointer(); } - DiskAnnSearcherEntity *entity = new (std::nothrow) DiskAnnSearcherEntity( - meta_header_, pq_meta_, meta_segment, pq_meta_segment, pq_data_segment, - vector_segment, key_segment, key_mapping_segment, entrypoint_segment, - num_threads_, list_size_, cache_nodes_num_, warm_up_, beam_size_, meta_, - pq_table_, key_buffer_, key_mapping_buffer_, entrypoints_); - if (ailego_unlikely(!entity)) { - LOG_ERROR("DiskAnnSearcherEntity new failed"); - } - - return DiskAnnEntity::Pointer(entity); + return DiskAnnEntity::Pointer(entity.release()); } int DiskAnnSearcherEntity::load(const IndexMeta &meta, @@ -158,7 +191,7 @@ int DiskAnnSearcherEntity::load_pq_segment() { // 1. read pq meta read_size = pq_meta_segment_->read(offset, &data, sizeof(DiskAnnPqMeta)); - if (read_size != sizeof(DiskAnnPqMeta)) { + if (read_size != sizeof(DiskAnnPqMeta) || data == nullptr) { LOG_ERROR("Read segment %s failed, expect: %zu, actual: %zu", DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str(), sizeof(DiskAnnPqMeta), read_size); @@ -168,57 +201,101 @@ int DiskAnnSearcherEntity::load_pq_segment() { memcpy(reinterpret_cast(&pq_meta_), data, sizeof(DiskAnnPqMeta)); offset += read_size; + uint32_t pq_dimension = 0; + size_t expected_full_pivot_size = 0; + size_t expected_centroid_size = 0; + int layout_ret = get_pq_layout( + meta_, &pq_dimension, &expected_full_pivot_size, &expected_centroid_size); + if (layout_ret != 0) { + LOG_ERROR("Invalid DiskAnn PQ layout for the configured index metadata"); + return layout_ret; + } + + size_t chunk_offsets_count = 0; + size_t expected_chunk_offsets_size = 0; if (pq_meta_.chunk_num == 0 || meta_header_.doc_cnt == 0 || - pq_meta_.full_pivot_data_size == 0 || pq_meta_.centroid_data_size == 0 || - pq_meta_.chunk_num > meta_.dimension()) { - LOG_ERROR("Invalid empty DiskAnn PQ metadata"); + pq_meta_.chunk_num > pq_dimension || + pq_meta_.full_pivot_data_size != expected_full_pivot_size || + pq_meta_.centroid_data_size != expected_centroid_size || + pq_meta_.chunk_num > std::numeric_limits::max() || + !checked_add(static_cast(pq_meta_.chunk_num), 1, + &chunk_offsets_count) || + !checked_multiply(chunk_offsets_count, sizeof(uint32_t), + &expected_chunk_offsets_size) || + (pq_meta_.chunk_offsets_size != 0 && + pq_meta_.chunk_offsets_size != expected_chunk_offsets_size)) { + LOG_ERROR("Invalid DiskAnn PQ metadata sizes"); + return IndexError_InvalidFormat; + } + + size_t pq_meta_data_size = sizeof(DiskAnnPqMeta); + if (!checked_add(pq_meta_data_size, expected_full_pivot_size, + &pq_meta_data_size) || + !checked_add(pq_meta_data_size, expected_centroid_size, + &pq_meta_data_size) || + !checked_add(pq_meta_data_size, expected_chunk_offsets_size, + &pq_meta_data_size) || + pq_meta_segment_->data_size() < pq_meta_data_size) { + LOG_ERROR("DiskAnn PQ metadata segment is shorter than its layout"); return IndexError_InvalidFormat; } // 2. read full pivot data std::vector full_pivot_data; - full_pivot_data.resize(pq_meta_.full_pivot_data_size); + std::vector centroid; + std::vector chunk_offsets; + try { + full_pivot_data.resize(expected_full_pivot_size); + centroid.resize(expected_centroid_size); + chunk_offsets.resize(chunk_offsets_count); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate DiskAnn PQ metadata buffers"); + return IndexError_NoMemory; + } - read_size = - pq_meta_segment_->read(offset, &data, pq_meta_.full_pivot_data_size); - if (read_size != pq_meta_.full_pivot_data_size) { + read_size = pq_meta_segment_->read(offset, &data, expected_full_pivot_size); + if (read_size != expected_full_pivot_size || data == nullptr) { LOG_ERROR("Read segment %s failed, expect: %zu, actual: %zu", DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str(), - (size_t)(pq_meta_.full_pivot_data_size), (size_t)read_size); + expected_full_pivot_size, read_size); return IndexError_ReadData; } - memcpy(&(full_pivot_data[0]), data, read_size); + memcpy(full_pivot_data.data(), data, read_size); offset += read_size; // 3. read centroid - std::vector centroid; - centroid.resize(pq_meta_.centroid_data_size); - - read_size = - pq_meta_segment_->read(offset, &data, pq_meta_.centroid_data_size); - if (read_size != pq_meta_.centroid_data_size) { + read_size = pq_meta_segment_->read(offset, &data, expected_centroid_size); + if (read_size != expected_centroid_size || data == nullptr) { LOG_ERROR("Read segment %s failed, expect: %zu, actual: %zu", DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str(), - (size_t)(pq_meta_.centroid_data_size), (size_t)read_size); + expected_centroid_size, read_size); return IndexError_ReadData; } - memcpy(&(centroid[0]), data, read_size); + memcpy(centroid.data(), data, read_size); offset += read_size; // 4. chunk offset - std::vector chunk_offsets; - chunk_offsets.resize(pq_meta_.chunk_num + 1); - - read_size = pq_meta_segment_->read( - offset, &data, (pq_meta_.chunk_num + 1) * sizeof(uint32_t)); - if (read_size != (pq_meta_.chunk_num + 1) * sizeof(uint32_t)) { + read_size = + pq_meta_segment_->read(offset, &data, expected_chunk_offsets_size); + if (read_size != expected_chunk_offsets_size || data == nullptr) { LOG_ERROR("Read segment %s failed, expect: %zu, actual: %zu", DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str(), - (size_t)((pq_meta_.chunk_num + 1) * sizeof(uint32_t)), - (size_t)read_size); + expected_chunk_offsets_size, read_size); return IndexError_ReadData; } - memcpy(&(chunk_offsets[0]), data, read_size); + memcpy(chunk_offsets.data(), data, read_size); + + if (chunk_offsets.front() != 0 || chunk_offsets.back() != pq_dimension) { + LOG_ERROR("Invalid DiskAnn PQ chunk offset boundaries"); + return IndexError_InvalidFormat; + } + for (size_t i = 1; i < chunk_offsets.size(); ++i) { + if (chunk_offsets[i - 1] >= chunk_offsets[i] || + chunk_offsets[i] > pq_dimension) { + LOG_ERROR("Invalid DiskAnn PQ chunk offsets"); + return IndexError_InvalidFormat; + } + } // load pq data std::vector pq_data; @@ -229,22 +306,40 @@ int DiskAnnSearcherEntity::load_pq_segment() { return IndexError_InvalidFormat; } - pq_data.resize(meta_header_.doc_cnt * pq_meta_.chunk_num); + size_t pq_data_size = 0; + if (meta_header_.doc_cnt > std::numeric_limits::max() || + !checked_multiply(static_cast(meta_header_.doc_cnt), + static_cast(pq_meta_.chunk_num), + &pq_data_size) || + pq_data_segment_->data_size() < pq_data_size) { + LOG_ERROR("Invalid DiskAnn PQ data size"); + return IndexError_InvalidFormat; + } + + try { + pq_data.resize(pq_data_size); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate DiskAnn PQ data buffer"); + return IndexError_NoMemory; + } - void *pq_data_ptr = &pq_data[0]; - read_size = pq_data_segment_->fetch( - 0, pq_data_ptr, meta_header_.doc_cnt * pq_meta_.chunk_num); + read_size = pq_data_segment_->fetch(0, pq_data.data(), pq_data_size); - if (read_size != meta_header_.doc_cnt * pq_meta_.chunk_num) { + if (read_size != pq_data_size) { LOG_ERROR("Read segment %s failed, expect: %zu, actual: %zu", - DiskAnnEntity::kDiskAnnPqMetaSegmentId.c_str(), - (size_t)(meta_header_.doc_cnt * pq_meta_.chunk_num), - (size_t)read_size); + DiskAnnEntity::kDiskAnnPqDataSegmentId.c_str(), pq_data_size, + read_size); return IndexError_ReadData; } - pq_table_ = std::make_shared(meta_, pq_meta_.chunk_num); + try { + pq_table_ = std::make_shared( + meta_, static_cast(pq_meta_.chunk_num)); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate DiskAnn PQ table"); + return IndexError_NoMemory; + } return pq_table_->init(full_pivot_data, centroid, chunk_offsets, pq_data); } @@ -259,13 +354,28 @@ int DiskAnnSearcherEntity::load_header_segment() { } if (meta_segment_->read(0, reinterpret_cast(&data), sizeof(DiskAnnMetaHeader)) != - sizeof(DiskAnnMetaHeader)) { + sizeof(DiskAnnMetaHeader) || + data == nullptr) { LOG_ERROR("Read segment %s failed", kDiskAnnMetaSegmentId.c_str()); return IndexError_ReadData; } memcpy(reinterpret_cast(&meta_header_), data, sizeof(DiskAnnMetaHeader)); + if (meta_header_.doc_cnt == 0 || meta_header_.doc_cnt > kInvalidId) { + LOG_ERROR("Invalid DiskAnn document count: %" PRIu64, meta_header_.doc_cnt); + return IndexError_InvalidFormat; + } + if (meta_header_.ndims != meta_.dimension()) { + LOG_ERROR("Invalid DiskAnn dimension: stored=%" PRIu64 " expected=%u", + meta_header_.ndims, meta_.dimension()); + return IndexError_InvalidFormat; + } + if (meta_header_.medoid >= meta_header_.doc_cnt) { + LOG_ERROR("Invalid DiskAnn medoid: %" PRIu64, meta_header_.medoid); + return IndexError_InvalidFormat; + } + return 0; } @@ -289,17 +399,32 @@ int DiskAnnSearcherEntity::load_key_segment() { return IndexError_InvalidFormat; } - size_t key_data_len = doc_cnt() * sizeof(diskann_key_t); + size_t key_data_len = 0; + if (doc_cnt() > std::numeric_limits::max() || + !checked_multiply(static_cast(doc_cnt()), sizeof(diskann_key_t), + &key_data_len) || + key_segment_->data_size() < key_data_len) { + LOG_ERROR("Invalid DiskAnn key segment size"); + return IndexError_InvalidFormat; + } const void *data = nullptr; if (key_segment_->read(0, reinterpret_cast(&data), - key_data_len) != key_data_len) { + key_data_len) != key_data_len || + data == nullptr) { LOG_ERROR("Read segment %s failed", kDiskAnnKeySegmentId.c_str()); return IndexError_ReadData; } - key_buffer_.resize(key_data_len); - memcpy(&(key_buffer_[0]), data, key_data_len); + try { + auto key_buffer = std::make_shared>( + static_cast(doc_cnt())); + memcpy(key_buffer->data(), data, key_data_len); + key_buffer_ = std::move(key_buffer); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate DiskAnn key buffer"); + return IndexError_NoMemory; + } return 0; } @@ -312,10 +437,15 @@ int DiskAnnSearcherEntity::load_entrypoint_segment() { return IndexError_InvalidFormat; } - const void *data = nullptr; + if (entrypoint_segment_->data_size() < sizeof(uint32_t)) { + LOG_ERROR("Invalid segment %s size", kDiskAnnEntryPointSegmentId.c_str()); + return IndexError_InvalidFormat; + } + const void *data = nullptr; if (entrypoint_segment_->read(0, reinterpret_cast(&data), - sizeof(uint32_t)) != sizeof(uint32_t)) { + sizeof(uint32_t)) != sizeof(uint32_t) || + data == nullptr) { LOG_ERROR("Read segment %s failed", kDiskAnnEntryPointSegmentId.c_str()); return IndexError_ReadData; } @@ -323,18 +453,53 @@ int DiskAnnSearcherEntity::load_entrypoint_segment() { uint32_t entrypoint_cnt = 0; memcpy(&entrypoint_cnt, data, sizeof(uint32_t)); - if (entrypoint_cnt != 0) { - size_t entrypoint_data_len = entrypoint_cnt * sizeof(diskann_id_t); + size_t entrypoint_data_len = 0; + size_t expected_segment_size = 0; + if (entrypoint_cnt > meta_header_.doc_cnt || + !checked_multiply(entrypoint_cnt, sizeof(diskann_id_t), + &entrypoint_data_len) || + !checked_add(sizeof(uint32_t), entrypoint_data_len, + &expected_segment_size) || + entrypoint_segment_->data_size() != expected_segment_size) { + LOG_ERROR("Invalid DiskAnn entrypoint count or segment size: count=%u", + entrypoint_cnt); + return IndexError_InvalidFormat; + } - if (entrypoint_segment_->read(sizeof(uint32_t), - reinterpret_cast(&data), - entrypoint_data_len) != entrypoint_data_len) { - LOG_ERROR("Read segment %s failed", kDiskAnnEntryPointSegmentId.c_str()); - return IndexError_ReadData; - } + std::vector entrypoints; + try { + entrypoints.resize(entrypoint_cnt); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate DiskAnn entrypoints"); + return IndexError_NoMemory; + } catch (const std::length_error &) { + LOG_ERROR("Invalid DiskAnn entrypoint count: %u", entrypoint_cnt); + return IndexError_InvalidFormat; + } - entrypoints_.resize(entrypoint_cnt); - memcpy(&(entrypoints_[0]), data, entrypoint_data_len); + if (entrypoint_data_len != 0 && + (entrypoint_segment_->read(sizeof(uint32_t), + reinterpret_cast(&data), + entrypoint_data_len) != entrypoint_data_len || + data == nullptr)) { + LOG_ERROR("Read segment %s failed", kDiskAnnEntryPointSegmentId.c_str()); + return IndexError_ReadData; + } + if (entrypoint_data_len != 0) { + memcpy(entrypoints.data(), data, entrypoint_data_len); + } + for (diskann_id_t id : entrypoints) { + if (id >= meta_header_.doc_cnt) { + LOG_ERROR("Invalid DiskAnn entrypoint id: %u", id); + return IndexError_InvalidFormat; + } + } + try { + entrypoints_ = std::make_shared>( + std::move(entrypoints)); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to retain DiskAnn entrypoints"); + return IndexError_NoMemory; } return 0; @@ -349,39 +514,90 @@ int DiskAnnSearcherEntity::load_key_mapping_segment() { return IndexError_InvalidFormat; } - size_t key_mapping_data_len = doc_cnt() * sizeof(diskann_id_t); + if (!key_buffer_ || key_buffer_->size() != doc_cnt()) { + LOG_ERROR("DiskAnn keys must be loaded before the key mapping"); + return IndexError_InvalidFormat; + } + + size_t key_mapping_data_len = 0; + if (doc_cnt() > std::numeric_limits::max() || + !checked_multiply(static_cast(doc_cnt()), sizeof(diskann_id_t), + &key_mapping_data_len) || + key_mapping_segment_->data_size() < key_mapping_data_len) { + LOG_ERROR("Invalid DiskAnn key mapping segment size"); + return IndexError_InvalidFormat; + } const void *data = nullptr; if (key_mapping_segment_->read(0, reinterpret_cast(&data), key_mapping_data_len) != - key_mapping_data_len) { + key_mapping_data_len || + data == nullptr) { LOG_ERROR("Read segment %s failed", kDiskAnnKeyMappingSegmentId.c_str()); return IndexError_ReadData; } - key_mapping_buffer_.resize(key_mapping_data_len); - memcpy(&(key_mapping_buffer_[0]), data, key_mapping_data_len); + try { + auto key_mapping_buffer = std::make_shared>( + static_cast(doc_cnt())); + memcpy(key_mapping_buffer->data(), data, key_mapping_data_len); + + std::vector seen_mapping_ids(key_mapping_buffer->size(), 0); + diskann_key_t previous_key = 0; + bool have_previous_key = false; + bool reached_invalid_keys = false; + for (size_t i = 0; i < key_mapping_buffer->size(); ++i) { + const diskann_id_t local_id = (*key_mapping_buffer)[i]; + if (local_id >= key_buffer_->size()) { + LOG_ERROR("Invalid DiskAnn key mapping id: %u", local_id); + return IndexError_InvalidFormat; + } + if (seen_mapping_ids[local_id] != 0) { + LOG_ERROR("Duplicate DiskAnn key mapping id: %u", local_id); + return IndexError_InvalidFormat; + } + seen_mapping_ids[local_id] = 1; + + const diskann_key_t local_key = (*key_buffer_)[local_id]; + if (local_key == kInvalidKey) { + reached_invalid_keys = true; + continue; + } + if (reached_invalid_keys || + (have_previous_key && local_key <= previous_key)) { + LOG_ERROR("DiskAnn key mapping is not strictly ordered"); + return IndexError_InvalidFormat; + } + previous_key = local_key; + have_previous_key = true; + } + key_mapping_buffer_ = std::move(key_mapping_buffer); + } catch (const std::bad_alloc &) { + LOG_ERROR("Failed to allocate DiskAnn key mapping buffer"); + return IndexError_NoMemory; + } return 0; } //! Get vector local id by key diskann_id_t DiskAnnSearcherEntity::get_id(diskann_key_t key) const { - const diskann_id_t *key_mapping_data_ptr = - reinterpret_cast(key_mapping_buffer_.data()); - - const diskann_key_t *key_data_ptr = - reinterpret_cast(key_buffer_.data()); + if (key == kInvalidKey || !key_mapping_buffer_ || !key_buffer_ || + key_mapping_buffer_->size() != key_buffer_->size()) { + return kInvalidId; + } //! Do binary search - diskann_id_t start = 0UL; - diskann_id_t end = doc_cnt(); - diskann_id_t idx = 0u; + size_t start = 0; + size_t end = key_mapping_buffer_->size(); while (start < end) { - idx = start + (end - start) / 2; - diskann_id_t local_id = key_mapping_data_ptr[idx]; + const size_t idx = start + (end - start) / 2; + const diskann_id_t local_id = (*key_mapping_buffer_)[idx]; + if (local_id >= key_buffer_->size()) { + return kInvalidId; + } - const diskann_key_t local_key = key_data_ptr[local_id]; + const diskann_key_t local_key = (*key_buffer_)[local_id]; if (local_key < key) { start = idx + 1; @@ -396,67 +612,10 @@ diskann_id_t DiskAnnSearcherEntity::get_id(diskann_key_t key) const { } diskann_key_t DiskAnnSearcherEntity::get_key(diskann_id_t id) const { - const diskann_key_t *key_data_ptr = - reinterpret_cast(key_buffer_.data()); - - return key_data_ptr[id]; -} - -const void *DiskAnnSearcherEntity::get_vector(diskann_id_t id) const { - if (!vector_segment_) { - LOG_ERROR("Vector segment is null"); - return nullptr; - } - - uint64_t sector_offset = - DiskAnnUtil::get_node_sector(node_per_sector(), max_node_size(), - DiskAnnUtil::kSectorSize, id) * - DiskAnnUtil::kSectorSize; - uint64_t within_sector_offset = - (node_per_sector() == 0 ? 0 : (id % node_per_sector()) * max_node_size()); - uint64_t total_offset = sector_offset + within_sector_offset; - - size_t read_size = meta_.element_size(); - 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: %llu", id, - (unsigned long long)total_offset); - return nullptr; + if (!key_buffer_ || id >= key_buffer_->size()) { + return kInvalidKey; } - - return vec; -} - -std::pair DiskAnnSearcherEntity::get_neighbors( - diskann_id_t id) const { - if (!vector_segment_) { - return std::make_pair(0, nullptr); - } - - uint64_t read_sector_offset = - DiskAnnUtil::get_node_sector(node_per_sector(), max_node_size(), - DiskAnnUtil::kSectorSize, id) * - DiskAnnUtil::kSectorSize; - uint64_t node_vec_offset = - read_sector_offset + - (node_per_sector() == 0 ? 0 : (id % node_per_sector()) * max_node_size()); - - const void *data; - if (ailego_unlikely( - vector_segment_->read(node_vec_offset, &data, max_node_size()) != - max_node_size())) { - LOG_ERROR("Read neighbors from segment failed"); - return {0, nullptr}; - } - - const uint8_t *data_ptr = reinterpret_cast(data); - const diskann_id_t *node_neighbor = - reinterpret_cast(data_ptr + meta_.element_size()); - - auto neighbor_num = *node_neighbor; - - return std::make_pair(neighbor_num, node_neighbor + 1); + return (*key_buffer_)[id]; } } // namespace core diff --git a/src/core/algorithm/diskann/diskann_searcher_entity.h b/src/core/algorithm/diskann/diskann_searcher_entity.h index 439b2b067..590a248e4 100644 --- a/src/core/algorithm/diskann/diskann_searcher_entity.h +++ b/src/core/algorithm/diskann/diskann_searcher_entity.h @@ -22,6 +22,8 @@ namespace zvec { namespace core { +class DiskAnnCacheTestPeer; + class DiskAnnSearcherEntity : public DiskAnnEntity { public: using Pointer = std::shared_ptr; @@ -35,6 +37,7 @@ class DiskAnnSearcherEntity : public DiskAnnEntity { const DiskAnnEntity::Pointer clone() const override; void clear(); + void release_storage(); int load(const IndexMeta &meta, IndexStorage::Pointer storage); int load_pq_segment(); int load_header_segment(); @@ -55,48 +58,15 @@ class DiskAnnSearcherEntity : public DiskAnnEntity { return vector_segment_; } - std::vector &entrypoints() { - return entrypoints_; + const std::vector &entrypoints() const { + static const std::vector empty; + return entrypoints_ == nullptr ? empty : *entrypoints_; } - std::pair get_neighbors( - diskann_id_t id) const override; - diskann_id_t get_id(diskann_key_t key) const override; diskann_key_t get_key(diskann_id_t id) const override; - const void *get_vector(diskann_id_t id) const override; private: - DiskAnnSearcherEntity( - const DiskAnnMetaHeader &meta_header, const DiskAnnPqMeta &pq_meta, - const SegmentPointer &meta_segment, const SegmentPointer &pq_meta_segment, - const SegmentPointer &pq_data_segment, - const SegmentPointer &vector_segment, const SegmentPointer &key_segment, - const SegmentPointer &key_mapping_segment, - const SegmentPointer &entrypoint_segment, uint32_t num_threads, - uint32_t list_size, uint32_t cache_nodes_num, bool warm_up, - uint32_t beam_size, const IndexMeta meta, PQTable::Pointer pq_table, - const std::string &key_buffer, const std::string &key_mapping_buffer, - const std::vector &entrypoints) - : DiskAnnEntity(meta_header, pq_meta), - meta_segment_(meta_segment), - pq_meta_segment_(pq_meta_segment), - pq_data_segment_(pq_data_segment), - vector_segment_(vector_segment), - key_segment_(key_segment), - key_mapping_segment_(key_mapping_segment), - entrypoint_segment_{entrypoint_segment}, - num_threads_{num_threads}, - list_size_{list_size}, - cache_nodes_num_{cache_nodes_num}, - warm_up_{warm_up}, - beam_size_{beam_size}, - meta_{meta}, - pq_table_{pq_table}, - key_buffer_{key_buffer}, - key_mapping_buffer_{key_mapping_buffer}, - entrypoints_{entrypoints} {} - IndexStorage::Pointer storage_{}; SegmentPointer meta_segment_{nullptr}; @@ -107,19 +77,14 @@ class DiskAnnSearcherEntity : public DiskAnnEntity { SegmentPointer key_mapping_segment_{nullptr}; SegmentPointer entrypoint_segment_{nullptr}; - uint32_t num_threads_{1}; - uint32_t list_size_{200}; - uint32_t cache_nodes_num_{0}; - - bool warm_up_{false}; - uint32_t beam_size_{2}; - IndexMeta meta_; PQTable::Pointer pq_table_; - std::string key_buffer_; - std::string key_mapping_buffer_; - std::vector entrypoints_; + std::shared_ptr> key_buffer_; + std::shared_ptr> key_mapping_buffer_; + std::shared_ptr> entrypoints_; + + friend class DiskAnnCacheTestPeer; }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_streamer.cc b/src/core/algorithm/diskann/diskann_streamer.cc index 1752161bd..00e1a6fac 100644 --- a/src/core/algorithm/diskann/diskann_streamer.cc +++ b/src/core/algorithm/diskann/diskann_streamer.cc @@ -13,6 +13,9 @@ // limitations under the License. #include "diskann_streamer.h" +#include +#include +#include #include "diskann_context.h" #include "diskann_index_provider.h" #include "diskann_indexer.h" @@ -47,15 +50,33 @@ int DiskAnnStreamer::init(const IndexMeta &meta, return IndexError_NoReady; } - meta_ = meta; - params_ = search_params; - list_size_ = 200; - cache_nodes_num_ = 0; - log_diskann_io_backend(); - params_.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size_); - params_.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, &cache_nodes_num_); + uint32_t list_size = 200; + uint32_t cache_nodes_num = 0; + search_params.get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size); + if (list_size == 0) { + LOG_ERROR("list_size must be positive"); + return IndexError_InvalidArgument; + } + long long configured_cache_nodes = 0; + if (search_params.get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, + &configured_cache_nodes)) { + if (configured_cache_nodes < 0 || + static_cast(configured_cache_nodes) > + std::numeric_limits::max()) { + LOG_ERROR("cache_node_num must be in [0, UINT32_MAX]"); + return IndexError_InvalidArgument; + } + cache_nodes_num = static_cast(configured_cache_nodes); + } + + // Commit only after every value has been validated. A failed re-init must + // leave either the previous valid configuration or STATE_INIT untouched. + meta_ = meta; + params_ = search_params; + list_size_ = list_size; + cache_nodes_num_ = cache_nodes_num; state_ = STATE_INITED; return 0; } @@ -87,6 +108,12 @@ int DiskAnnStreamer::open(IndexStorage::Pointer storage) { LOG_ERROR("Initialize and close DiskAnnStreamer before opening an index"); return IndexError_NoReady; } + if (!storage->file()) { + LOG_ERROR( + "DiskAnn requires storage with a shared file handle; disable " + "proxima.file.read_storage.alone_file_handle"); + return IndexError_InvalidArgument; + } { std::lock_guard lock(fetch_mutex_); @@ -119,19 +146,9 @@ int DiskAnnStreamer::open(IndexStorage::Pointer storage) { return res; } - if (cache_nodes_num_ != 0) { - std::vector node_list; - LOG_INFO("Caching %u nodes around medoid(s)", cache_nodes_num_); - - diskann_indexer_->cache_bfs_levels(cache_nodes_num_, node_list); - - ret = diskann_indexer_->load_cache_list(node_list); - if (ret != 0) { - return ret; - } - - node_list.clear(); - node_list.shrink_to_fit(); + ret = diskann_indexer_->configure_cache(cache_nodes_num_); + if (ret != 0) { + return ret; } measure_ = IndexFactory::CreateMetric(meta_.metric_name()); @@ -238,6 +255,7 @@ int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, if (ret != 0) { return ret; } + AILEGO_DEFER(diskann_indexer_.get(), &DiskAnnIndexer::release_io_ctx, ctx); if (ailego_unlikely(!group_options_valid(ctx))) { LOG_ERROR("Group search requires a callback and a positive group topk"); return IndexError_InvalidArgument; @@ -247,6 +265,7 @@ int DiskAnnStreamer::search_impl(const void *query, const IndexQueryMeta &qmeta, ctx->resize_results(count); for (uint32_t i = 0; i < count; i++) { + ctx->visit_filter().clear(); ctx->reset_query(query); ret = diskann_indexer_->knn_search(ctx); @@ -292,6 +311,7 @@ int DiskAnnStreamer::search_bf_impl(const void *query, if (ret != 0) { return ret; } + AILEGO_DEFER(diskann_indexer_.get(), &DiskAnnIndexer::release_io_ctx, ctx); if (ailego_unlikely(!group_options_valid(ctx))) { LOG_ERROR("Group search requires a callback and a positive group topk"); return IndexError_InvalidArgument; @@ -352,6 +372,7 @@ int DiskAnnStreamer::search_bf_by_p_keys_impl( if (ret != 0) { return ret; } + AILEGO_DEFER(diskann_indexer_.get(), &DiskAnnIndexer::release_io_ctx, ctx); if (ailego_unlikely(!group_options_valid(ctx))) { LOG_ERROR("Group search requires a callback and a positive group topk"); return IndexError_InvalidArgument; @@ -429,21 +450,13 @@ int DiskAnnStreamer::get_vector_by_id(const uint32_t id, std::lock_guard lock(fetch_mutex_); if (!fetch_ctx_) { - fetch_ctx_ = create_context(); + const DiskAnnEntity::Pointer fetch_entity = entity_.clone(); + fetch_ctx_ = + DiskAnnContext::create_fetch_context(meta_, measure_, fetch_entity); if (!fetch_ctx_) { - LOG_ERROR("Failed to create context for get_vector_by_id"); + LOG_ERROR("Failed to create fetch context for get_vector_by_id"); return IndexError_Runtime; } - } else { - auto *ctx = dynamic_cast(fetch_ctx_.get()); - if (!ctx) { - LOG_ERROR("Cast fetch context to DiskAnnContext failed"); - return IndexError_Cast; - } - int ret = ensure_compatible_context(fetch_ctx_, ctx); - if (ret != 0) { - return ret; - } } int ret = diskann_indexer_->get_vector(id, fetch_ctx_, fetch_vector_buffer_); @@ -473,8 +486,14 @@ IndexSearcher::Provider::Pointer DiskAnnStreamer::create_provider(void) const { LOG_ERROR("Failed to clone DiskAnn entity for provider"); return nullptr; } - return IndexProvider::Pointer(new (std::nothrow) DiskAnnIndexProvider( - meta_, entity, "DiskAnnStreamer")); + std::unique_ptr provider( + new (std::nothrow) DiskAnnIndexProvider( + meta_, measure_, entity, diskann_indexer_, "DiskAnnStreamer")); + if (!provider || !provider->ready()) { + LOG_ERROR("Failed to initialize DiskAnn provider dependencies"); + return nullptr; + } + return IndexProvider::Pointer(provider.release()); } IndexSearcher::Context::Pointer DiskAnnStreamer::create_context() const { diff --git a/src/core/algorithm/diskann/diskann_streamer.h b/src/core/algorithm/diskann/diskann_streamer.h index 43c36712a..ae713e52e 100644 --- a/src/core/algorithm/diskann/diskann_streamer.h +++ b/src/core/algorithm/diskann/diskann_streamer.h @@ -18,11 +18,11 @@ #include "diskann_context.h" #include "diskann_indexer.h" -class LinuxAlignedFileReader; - namespace zvec { namespace core { +class DiskAnnStreamerTestPeer; + class DiskAnnStreamer : public IndexStreamer { public: using ContextPointer = IndexStreamer::Context::Pointer; @@ -116,7 +116,7 @@ class DiskAnnStreamer : public IndexStreamer { //! Create a searcher context ContextPointer create_context() const override; - //! Create a vector iterator backed by the on-disk vector segment. + //! Create a vector iterator backed by the aligned DiskAnn file reader. //! Used by the merge code path (``MixedStreamerReducer``) to walk every //! vector held by this streamer. IndexSearcher::Provider::Pointer create_provider(void) const override; @@ -163,13 +163,11 @@ class DiskAnnStreamer : public IndexStreamer { uint32_t list_size_{200}; uint32_t cache_nodes_num_{0}; - bool warm_up_{false}; - uint32_t beam_size_{2}; DiskAnnIndexer::Pointer diskann_indexer_{nullptr}; DiskAnnSearcherEntity entity_{}; - // Fetches share the expensive I/O context, while returned MemoryBlocks own + // Fetches share a lightweight I/O context, while returned MemoryBlocks own // independent copies so their lifetime does not depend on this buffer. mutable std::mutex fetch_mutex_; mutable ContextPointer fetch_ctx_{}; @@ -179,6 +177,8 @@ class DiskAnnStreamer : public IndexStreamer { Stats stats_; State state_{STATE_INIT}; + + friend class DiskAnnStreamerTestPeer; }; } // namespace core diff --git a/src/core/algorithm/diskann/diskann_util.h b/src/core/algorithm/diskann/diskann_util.h index c22268ae1..c348db195 100644 --- a/src/core/algorithm/diskann/diskann_util.h +++ b/src/core/algorithm/diskann/diskann_util.h @@ -13,10 +13,15 @@ // limitations under the License. #pragma once +#include #include #include #include "diskann_entity.h" +#if defined(_WIN32) || defined(_WIN64) +#include +#endif + namespace zvec { namespace core { @@ -25,6 +30,13 @@ class DiskAnnUtil { static constexpr uint64_t kSectorSize = 4096; static constexpr uint64_t kMaxSectorReadNum = 128; + static constexpr uint64_t cache_load_batch_size( + uint64_t sector_num_per_node) { + return sector_num_per_node == 0 || sector_num_per_node > kMaxSectorReadNum + ? 1 + : kMaxSectorReadNum / sector_num_per_node; + } + public: static inline size_t div_round_up(size_t x, size_t y) { return (x / y + (x % y != 0)); @@ -35,23 +47,33 @@ class DiskAnnUtil { } static inline void alloc_aligned(void **ptr, size_t size, size_t align) { + if (ptr == nullptr) { + return; + } if (size == 0) { *ptr = nullptr; return; } +#if defined(_WIN32) || defined(_WIN64) + *ptr = ::_aligned_malloc(size, align); +#else // Unlike aligned_alloc(), posix_memalign() does not require size to be an // integral multiple of alignment and is available on Linux and macOS. if (::posix_memalign(ptr, align, size) != 0) { *ptr = nullptr; } +#endif } static inline void free_aligned(void *ptr) { if (ptr == nullptr) { return; } - +#if defined(_WIN32) || defined(_WIN64) + ::_aligned_free(ptr); +#else free(ptr); +#endif } template @@ -86,9 +108,14 @@ class DiskAnnUtil { : node_id * div_round_up(max_nodesize_, sectorsize_)); } - static inline uint32_t *offset_to_node_neighbor(uint8_t *node_buf, - uint32_t elementsize_) { - return (uint32_t *)(node_buf + elementsize_); + static inline uint8_t *offset_to_node_neighbor(uint8_t *node_buf, + uint32_t elementsize_) { + return node_buf + elementsize_; + } + + static inline const uint8_t *offset_to_node_neighbor(const uint8_t *node_buf, + uint32_t elementsize_) { + return node_buf + elementsize_; } static inline uint8_t *offset_to_node(uint32_t node_per_sector, @@ -142,6 +169,9 @@ class NeighborPriorityQueue { : size_(0), capacity_(capacity), cur_(0), data_(capacity + 1) {} void insert(const Neighbor &nbr) { + if (capacity_ == 0) { + return; + } if (size_ == capacity_ && data_[size_ - 1] < nbr) { return; } diff --git a/src/core/framework/index_meta.cc b/src/core/framework/index_meta.cc index 8d98aadfd..e97faad6d 100644 --- a/src/core/framework/index_meta.cc +++ b/src/core/framework/index_meta.cc @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include #include #include @@ -37,6 +39,57 @@ struct IndexMetaFormatHeader { static_assert(sizeof(IndexMetaFormatHeader) % 32 == 0, "IndexMetaBufferFormat must be aligned with 32 bytes"); +namespace { + +bool ComputeElementSize(uint32_t data_type, uint32_t unit_size, + uint32_t dimension, uint32_t extra_meta_size, + uint32_t *element_size) { + if (data_type > static_cast(IndexMeta::DataType::DT_BINARY64)) { + return false; + } + + const auto type = static_cast(data_type); + const uint32_t expected_unit_size = IndexMeta::UnitSizeof(type); + if (unit_size != expected_unit_size) { + return false; + } + + uint64_t base_size = 0; + switch (type) { + case IndexMeta::DataType::DT_UNDEFINED: + break; + case IndexMeta::DataType::DT_FP16: + case IndexMeta::DataType::DT_FP32: + case IndexMeta::DataType::DT_FP64: + case IndexMeta::DataType::DT_INT8: + case IndexMeta::DataType::DT_INT16: + base_size = static_cast(dimension) * unit_size; + break; + case IndexMeta::DataType::DT_INT4: { + const uint64_t values_per_unit = static_cast(unit_size) * 2; + base_size = (static_cast(dimension) + values_per_unit - 1) / + values_per_unit * unit_size; + break; + } + case IndexMeta::DataType::DT_BINARY32: + case IndexMeta::DataType::DT_BINARY64: { + const uint64_t values_per_unit = static_cast(unit_size) * 8; + base_size = (static_cast(dimension) + values_per_unit - 1) / + values_per_unit * unit_size; + break; + } + } + + const uint64_t total_size = base_size + extra_meta_size; + if (total_size > std::numeric_limits::max()) { + return false; + } + *element_size = static_cast(total_size); + return true; +} + +} // namespace + void IndexMeta::serialize(std::string *out) const { ailego::Params attachment; IndexMetaFormatHeader format; @@ -119,56 +172,71 @@ void IndexMeta::serialize(std::string *out) const { attachment.set("attributes", attributes_); } - out->assign(reinterpret_cast(&format), sizeof(format)); - size_t offset = static_cast(out->size()); - + std::string attachment_buffer; if (!attachment.empty()) { - std::string buf; - ailego::Params::SerializeToBuffer(attachment, &buf); - out->append(buf.data(), buf.size()); - IndexMetaFormatHeader *header = (IndexMetaFormatHeader *)out->data(); - header->attachment_offset = static_cast(offset); - header->attachment_size = static_cast(buf.size()); - offset += buf.size(); + ailego::Params::SerializeToBuffer(attachment, &attachment_buffer); + if (attachment_buffer.size() > std::numeric_limits::max()) { + out->clear(); + return; + } + format.attachment_offset = sizeof(format); + format.attachment_size = static_cast(attachment_buffer.size()); } + out->assign(reinterpret_cast(&format), sizeof(format)); + out->append(attachment_buffer); } bool IndexMeta::deserialize(const void *data, size_t len) { - const IndexMetaFormatHeader *format = - reinterpret_cast(data); - this->clear(); - if (sizeof(IndexMetaFormatHeader) > len) { + if (data == nullptr || sizeof(IndexMetaFormatHeader) > len) { return false; } - if (sizeof(IndexMetaFormatHeader) > format->header_size) { + + IndexMetaFormatHeader format; + std::memcpy(&format, data, sizeof(format)); + if (format.header_size < sizeof(IndexMetaFormatHeader) || + format.header_size > len) { return false; } - meta_type_ = static_cast(format->meta_type); - major_order_ = static_cast(format->major_order); - data_type_ = static_cast(format->data_type); - dimension_ = format->dimension; - unit_size_ = format->unit_size; - extra_meta_size_ = format->extra_meta_size; - element_size_ = IndexMeta::ElementSizeof(data_type_, unit_size_, dimension_) + - extra_meta_size_; - space_id_ = format->space_id; + if (format.meta_type > + static_cast(IndexMeta::MetaType::MT_SPARSE) || + format.major_order > + static_cast(IndexMeta::MajorOrder::MO_COLUMN)) { + return false; + } + + uint32_t element_size = 0; + if (!ComputeElementSize(format.data_type, format.unit_size, format.dimension, + format.extra_meta_size, &element_size)) { + return false; + } // Read attachment ailego::Params attachment; - if (format->attachment_size) { - if (format->attachment_offset + format->attachment_size > len) { + if (format.attachment_size != 0) { + if (format.attachment_offset < format.header_size || + format.attachment_offset > len || + format.attachment_size > len - format.attachment_offset) { return false; } std::string str( - reinterpret_cast(data) + format->attachment_offset, - format->attachment_size); + reinterpret_cast(data) + format.attachment_offset, + format.attachment_size); if (!ailego::Params::ParseFromBuffer(str, &attachment)) { return false; } } + meta_type_ = static_cast(format.meta_type); + major_order_ = static_cast(format.major_order); + data_type_ = static_cast(format.data_type); + dimension_ = format.dimension; + unit_size_ = format.unit_size; + extra_meta_size_ = format.extra_meta_size; + element_size_ = element_size; + space_id_ = format.space_id; + ailego::Params item; if (attachment.get("metric", &item)) { item.get("name", &metric_name_); diff --git a/src/core/interface/index.cc b/src/core/interface/index.cc index 69e7d7c53..66f66a579 100644 --- a/src/core/interface/index.cc +++ b/src/core/interface/index.cc @@ -40,8 +40,9 @@ bool Index::init_context() { context_index_ = (magic_enum::enum_integer(param_.index_type) - 1) * 2 + static_cast(is_sparse_); if (_context_list[context_index_] == nullptr) { - if ((_context_list[context_index_] = streamer_->create_context()) == - nullptr) { + const auto streamer = streamer_snapshot(); + if (streamer == nullptr || (_context_list[context_index_] = + streamer->create_context()) == nullptr) { LOG_ERROR("Failed to create context"); return false; } @@ -64,25 +65,29 @@ BaseIndexParam::Pointer Index::get_param() const { } bool Index::is_trained() const { - return is_trained_; + return is_trained_.load(std::memory_order_acquire); } uint32_t Index::get_doc_count() const { - if (streamer_ == nullptr) { - return -1; + const auto streamer = streamer_snapshot(); + if (streamer == nullptr) { + return 0; } if (is_sparse_) { - return streamer_->create_sparse_provider()->count(); + const auto provider = streamer->create_sparse_provider(); + return provider == nullptr ? 0 : provider->count(); } - return streamer_->create_provider()->count(); + const auto provider = streamer->create_provider(); + return provider == nullptr ? 0 : provider->count(); } core::IndexStreamer::Pointer Index::index_searcher() { - return streamer_; + return streamer_snapshot(); } core::IndexProvider::Pointer Index::create_index_provider() const { - return streamer_->create_provider(); + const auto streamer = streamer_snapshot(); + return streamer == nullptr ? nullptr : streamer->create_provider(); } int Index::ParseMetricName(const BaseIndexParam ¶m) { @@ -271,6 +276,18 @@ int Index::CreateAndInitConverterReformer(const QuantizerParam ¶m, } int Index::Init(const BaseIndexParam ¶m) { + const int data_type = static_cast(param.data_type); + if (data_type <= static_cast(DataType::DT_UNDEFINED) || + data_type > static_cast(DataType::DT_BINARY64)) { + LOG_ERROR("Invalid data type: %d", data_type); + return core::IndexError_InvalidArgument; + } + if (param.dimension < 0 || param.dimension > MAX_DIMENSION || + (!param.is_sparse && param.dimension == 0)) { + LOG_ERROR("Invalid dimension: %d", param.dimension); + return core::IndexError_InvalidArgument; + } + param_ = param; // will lose the original type info is_sparse_ = param.is_sparse; @@ -372,7 +389,8 @@ int Index::open(const std::string &file_path, StorageOptions storage_options) { core::IndexError::What(ret)); return core::IndexError_Runtime; } - if (streamer_ == nullptr || streamer_->open(storage_) != 0) { + const auto streamer = streamer_snapshot(); + if (streamer == nullptr || streamer->open(storage_) != 0) { LOG_ERROR("Failed to open streamer, path: %s", file_path.c_str()); return core::IndexError_Runtime; } @@ -382,7 +400,7 @@ int Index::open(const std::string &file_path, StorageOptions storage_options) { // persisted meta loaded by the streamer. When there is no converter // (QuantizerType::kNone), reformer_ is nullptr by design. if (converter_ != nullptr && reformer_ == nullptr) { - const auto &meta = streamer_->meta(); + const auto &meta = streamer->meta(); if (meta.reformer_name().empty()) { LOG_ERROR( "Index::open: converter exists but reformer not initialized and " @@ -442,11 +460,13 @@ int Index::close() { return core::IndexError_Runtime; } } - if (ailego_unlikely(streamer_->cleanup() != 0)) { + const auto streamer = streamer_snapshot(); + if (streamer == nullptr || ailego_unlikely(streamer->cleanup() != 0)) { LOG_ERROR("Failed to cleanup streamer"); return core::IndexError_Runtime; } - if (ailego_unlikely(storage_->close() != 0)) { + const auto storage = storage_snapshot(); + if (storage == nullptr || ailego_unlikely(storage->close() != 0)) { LOG_ERROR("Failed to close storage"); return core::IndexError_Runtime; } @@ -464,11 +484,13 @@ int Index::flush() { LOG_ERROR("Cannot flush read-only index"); return core::IndexError_Runtime; } - if (ailego_unlikely(streamer_->flush(0) != 0)) { + const auto streamer = streamer_snapshot(); + if (streamer == nullptr || ailego_unlikely(streamer->flush(0) != 0)) { LOG_ERROR("Failed to flush streamer"); return core::IndexError_Runtime; } - if (ailego_unlikely(storage_->flush() != 0)) { + const auto storage = storage_snapshot(); + if (storage == nullptr || ailego_unlikely(storage->flush() != 0)) { LOG_ERROR("Failed to flush storage"); return core::IndexError_Runtime; } @@ -476,10 +498,11 @@ int Index::flush() { } bool Index::is_dirty() const { - if (!storage_) { + const auto storage = storage_snapshot(); + if (!storage) { return false; } - return storage_->is_dirty(); + return storage->is_dirty(); } int Index::fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { @@ -487,6 +510,10 @@ int Index::fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { LOG_ERROR("Index is not open"); return core::IndexError_Runtime; } + if (vector_data_buffer == nullptr) { + LOG_ERROR("Invalid output vector buffer"); + return core::IndexError_InvalidArgument; + } if (is_sparse_) { return _sparse_fetch(doc_id, vector_data_buffer); } @@ -627,8 +654,12 @@ int Index::search(const VectorData &vector_data, int Index::_dense_fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { + const auto streamer = streamer_snapshot(); + if (streamer == nullptr) { + return core::IndexError_NoReady; + } core::IndexStorage::MemoryBlock vector_block; - int ret = streamer_->get_vector_by_id(doc_id, vector_block); + int ret = streamer->get_vector_by_id(doc_id, vector_block); if (ret != 0) { LOG_ERROR("Failed to fetch vector, doc_id: %u", doc_id); return core::IndexError_Runtime; @@ -658,9 +689,13 @@ int Index::_dense_fetch(const uint32_t doc_id, int Index::_sparse_fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { + const auto streamer = streamer_snapshot(); + if (streamer == nullptr) { + return core::IndexError_NoReady; + } SparseVectorBuffer sparse_vector_buffer; - if (0 != streamer_->get_sparse_vector_by_id( + if (0 != streamer->get_sparse_vector_by_id( doc_id, &sparse_vector_buffer.count, &sparse_vector_buffer.indices, &sparse_vector_buffer.values)) { LOG_ERROR("Failed to fetch vector"); @@ -684,6 +719,10 @@ int Index::_sparse_fetch(const uint32_t doc_id, int Index::_dense_add(const VectorData &vector_data, const uint32_t doc_id, core::IndexContext::Pointer &context) { + const auto streamer = streamer_snapshot(); + if (streamer == nullptr) { + return core::IndexError_NoReady; + } if (!std::holds_alternative(vector_data.vector)) { LOG_ERROR("Invalid vector data"); return core::IndexError_Runtime; @@ -699,15 +738,15 @@ int Index::_dense_add(const VectorData &vector_data, const uint32_t doc_id, LOG_ERROR("Failed to convert vector"); return core::IndexError_Runtime; } - ret = streamer_->add_with_id_impl(doc_id, new_vector.data(), new_meta, - context); + ret = streamer->add_with_id_impl(doc_id, new_vector.data(), new_meta, + context); if (ret != 0) { LOG_ERROR("Failed to add vector"); return core::IndexError_Runtime; } } else { - int ret = streamer_->add_with_id_impl(doc_id, dense_vector.data, - input_vector_meta_, context); + int ret = streamer->add_with_id_impl(doc_id, dense_vector.data, + input_vector_meta_, context); if (ret != 0) { LOG_ERROR("Failed to add vector"); return core::IndexError_Runtime; @@ -719,6 +758,10 @@ int Index::_dense_add(const VectorData &vector_data, const uint32_t doc_id, int Index::_sparse_add(const VectorData &vector_data, const uint32_t doc_id, core::IndexContext::Pointer &context) { + const auto streamer = streamer_snapshot(); + if (streamer == nullptr) { + return core::IndexError_NoReady; + } if (!std::holds_alternative(vector_data.vector)) { LOG_ERROR("Invalid vector data"); return core::IndexError_Runtime; @@ -737,7 +780,7 @@ int Index::_sparse_add(const VectorData &vector_data, const uint32_t doc_id, LOG_ERROR("Failed to convert vector"); return core::IndexError_Runtime; } - ret = streamer_->add_with_id_impl( + ret = streamer->add_with_id_impl( doc_id, sparse_vector.count, sparse_vector.get_indices(), converted_sparse_values_buffer.data(), new_meta, context); if (ret != 0) { @@ -745,7 +788,7 @@ int Index::_sparse_add(const VectorData &vector_data, const uint32_t doc_id, return core::IndexError_Runtime; } } else { - int ret = streamer_->add_with_id_impl( + int ret = streamer->add_with_id_impl( doc_id, sparse_vector.count, sparse_vector.get_indices(), sparse_vector.get_values(), input_vector_meta_, context); if (ret != 0) { @@ -778,21 +821,25 @@ int Index::_dense_search(const VectorData &vector_data, } vector = new_vector.data(); } + const auto streamer = streamer_snapshot(); + if (streamer == nullptr) { + return core::IndexError_NoReady; + } if (search_param->bf_pks != nullptr) { // should we eliminate the copy of bf_pks? - if (streamer_->search_bf_by_p_keys_impl( + if (streamer->search_bf_by_p_keys_impl( vector, std::vector>{*search_param->bf_pks}, new_meta, 1, context) != 0) { LOG_ERROR("Failed to search_bf_by_p_keys_impl vector"); return core::IndexError_Runtime; } } else if (search_param->is_linear) { - if (streamer_->search_bf_impl(vector, new_meta, 1, context) != 0) { + if (streamer->search_bf_impl(vector, new_meta, 1, context) != 0) { LOG_ERROR("Failed to search vector"); return core::IndexError_Runtime; } } else { - if (streamer_->search_impl(vector, new_meta, 1, context) != 0) { + if (streamer->search_impl(vector, new_meta, 1, context) != 0) { LOG_ERROR("Failed to search vector"); return core::IndexError_Runtime; } @@ -907,8 +954,12 @@ int Index::_sparse_search(const VectorData &vector_data, values = converted_sparse_values_buffer.data(); } + const auto streamer = streamer_snapshot(); + if (streamer == nullptr) { + return core::IndexError_NoReady; + } if (search_param->bf_pks != nullptr) { - if (streamer_->search_bf_by_p_keys_impl( + if (streamer->search_bf_by_p_keys_impl( sparse_vector.count, indices, values, std::vector>{*search_param->bf_pks}, new_meta, context) != 0) { @@ -916,14 +967,14 @@ int Index::_sparse_search(const VectorData &vector_data, return core::IndexError_Runtime; } } else if (search_param->is_linear) { - if (streamer_->search_bf_impl(sparse_vector.count, indices, values, - new_meta, context) != 0) { + if (streamer->search_bf_impl(sparse_vector.count, indices, values, new_meta, + context) != 0) { LOG_ERROR("Failed to search vector"); return core::IndexError_Runtime; } } else { - if (streamer_->search_impl(sparse_vector.count, indices, values, new_meta, - context) != 0) { + if (streamer->search_impl(sparse_vector.count, indices, values, new_meta, + context) != 0) { LOG_ERROR("Failed to search vector"); return core::IndexError_Runtime; } @@ -1043,15 +1094,17 @@ int Index::merge(const std::vector &indexes, LOG_ERROR("Failed to init reducer"); return core::IndexError_Runtime; } - if (reducer->set_target_streamer_wiht_info(builder_, streamer_, converter_, - reformer_, + const auto target_streamer = streamer_snapshot(); + if (reducer->set_target_streamer_wiht_info(builder_, target_streamer, + converter_, reformer_, input_vector_meta_) != 0) { LOG_ERROR("Failed to set target streamer"); return core::IndexError_Runtime; } for (const auto &index : indexes) { - if (reducer->feed_streamer_with_reformer(index->streamer_, + const auto source_streamer = index->streamer_snapshot(); + if (reducer->feed_streamer_with_reformer(source_streamer, index->reformer_) != 0) { LOG_ERROR("Failed to feed streamer"); return core::IndexError_Runtime; diff --git a/src/core/interface/index_factory.cc b/src/core/interface/index_factory.cc index c640d5511..8e20436e1 100644 --- a/src/core/interface/index_factory.cc +++ b/src/core/interface/index_factory.cc @@ -132,6 +132,14 @@ BaseIndexParam::Pointer IndexFactory::DeserializeIndexParamFromJson( } return param; } + case IndexType::kDiskAnn: { + DiskAnnIndexParam::Pointer param = std::make_shared(); + if (!param->deserialize_from_json(json_str)) { + LOG_ERROR("Failed to deserialize diskann index param"); + return nullptr; + } + return param; + } case IndexType::kVamana: { VamanaIndexParam::Pointer param = std::make_shared(); if (!param->deserialize_from_json(json_str)) { @@ -203,6 +211,11 @@ std::string IndexFactory::QueryParamSerializeToJson(const QueryParamType ¶m, json_obj.set("nprobe", ailego::JsonValue(param.nprobe)); } index_type = IndexType::kIVFRabitq; + } else if constexpr (std::is_same_v) { + if (!omit_empty_value || param.list_size != 0) { + json_obj.set("list_size", ailego::JsonValue(param.list_size)); + } + index_type = IndexType::kDiskAnn; } else if constexpr (std::is_same_v) { if (!omit_empty_value || param.ef_search != 0) { json_obj.set("ef_search", ailego::JsonValue(param.ef_search)); @@ -231,6 +244,8 @@ template std::string IndexFactory::QueryParamSerializeToJson( const HNSWQueryParam ¶m, bool omit_empty_value); template std::string IndexFactory::QueryParamSerializeToJson( const IVFQueryParam ¶m, bool omit_empty_value); +template std::string IndexFactory::QueryParamSerializeToJson( + const DiskAnnQueryParam ¶m, bool omit_empty_value); template (); + if (!parse_common_fields(param)) { + return nullptr; + } + if (!extract_value_from_json(json_obj, "list_size", param->list_size, + tmp_json_value) || + param->list_size == 0) { + LOG_ERROR("Failed to deserialize DiskAnn list_size"); + return nullptr; + } + return param; } else if (index_type == IndexType::kVamana) { auto param = std::make_shared(); if (!parse_common_fields(param)) { @@ -406,6 +433,14 @@ typename QueryParamType::Pointer IndexFactory::QueryParamDeserializeFromJson( LOG_ERROR("Failed to deserialize nprobe"); return nullptr; } + } else if constexpr (std::is_same_v) { + if (index_type != IndexType::kDiskAnn || + !extract_value_from_json(json_obj, "list_size", param->list_size, + tmp_json_value) || + param->list_size == 0) { + LOG_ERROR("Failed to deserialize DiskAnn list_size"); + return nullptr; + } } else if constexpr (std::is_same_v) { if (!extract_value_from_json(json_obj, "ef_search", param->ef_search, tmp_json_value)) { @@ -440,6 +475,8 @@ template HNSWQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< HNSWQueryParam>(const std::string &json_str); template IVFQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< IVFQueryParam>(const std::string &json_str); +template DiskAnnQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< + DiskAnnQueryParam>(const std::string &json_str); template std::string IndexFactory::QueryParamSerializeToJson( const VamanaQueryParam ¶m, bool omit_empty_value); template VamanaQueryParam::Pointer IndexFactory::QueryParamDeserializeFromJson< diff --git a/src/core/interface/index_param.cc b/src/core/interface/index_param.cc index 1e75e682e..1349b714a 100644 --- a/src/core/interface/index_param.cc +++ b/src/core/interface/index_param.cc @@ -422,12 +422,27 @@ bool DiskAnnIndexParam::DeserializeFromJsonObject( return false; } + DESERIALIZE_VALUE_FIELD(json_obj, max_degree); + DESERIALIZE_VALUE_FIELD(json_obj, list_size); + DESERIALIZE_VALUE_FIELD(json_obj, pq_chunk_num); + + if (max_degree <= 0 || list_size <= 0 || pq_chunk_num < 0) { + LOG_ERROR( + "Invalid DiskAnn parameters: max_degree=%d list_size=%d " + "pq_chunk_num=%d", + max_degree, list_size, pq_chunk_num); + return false; + } + return true; } ailego::JsonObject DiskAnnIndexParam::SerializeToJsonObject( bool omit_empty_value) const { auto json_obj = BaseIndexParam::SerializeToJsonObject(omit_empty_value); + json_obj.set("max_degree", ailego::JsonValue(max_degree)); + json_obj.set("list_size", ailego::JsonValue(list_size)); + json_obj.set("pq_chunk_num", ailego::JsonValue(pq_chunk_num)); return json_obj; } diff --git a/src/core/interface/indexes/diskann_index.cc b/src/core/interface/indexes/diskann_index.cc index 493cc005b..7e1192ea1 100644 --- a/src/core/interface/indexes/diskann_index.cc +++ b/src/core/interface/indexes/diskann_index.cc @@ -12,9 +12,28 @@ // 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 +#if defined(_WIN32) || defined(_WIN64) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#else +#include +#include +#include +#endif #include #if DISKANN_SUPPORTED #include "algorithm/diskann/diskann_params.h" @@ -23,6 +42,172 @@ namespace zvec::core_interface { +#if DISKANN_SUPPORTED +namespace { + +std::string MakeSnapshotTemporaryPath(const std::string &file_path) { + static std::atomic sequence{0}; + const auto timestamp = + std::chrono::steady_clock::now().time_since_epoch().count(); + return file_path + ".merge-" + std::to_string(timestamp) + "-" + + std::to_string(sequence.fetch_add(1, std::memory_order_relaxed)) + + ".tmp"; +} + +enum class SnapshotReplaceResult { + kNotReplaced, + kReplaced, + kReplacedNotDurable, +}; + +#if defined(_WIN32) || defined(_WIN64) +bool ReplaceOpenFileWithPosixSemantics( + const std::filesystem::path &source_path, + const std::filesystem::path &destination_path) { + // FileRenameInfoEx is available at runtime on supported Windows Server + // versions, but the project still builds against an older SDK view. Keep + // the ABI declarations local so an open destination can be replaced while + // existing readers retain the old file object. + constexpr auto kFileRenameInfoEx = static_cast(22); + constexpr DWORD kReplaceIfExists = 0x00000001; + constexpr DWORD kPosixSemantics = 0x00000002; + struct ExtendedFileRenameInfo { + DWORD flags; + HANDLE root_directory; + DWORD file_name_length; + WCHAR file_name[1]; + }; + static_assert(offsetof(ExtendedFileRenameInfo, root_directory) == + offsetof(FILE_RENAME_INFO, RootDirectory)); + static_assert(offsetof(ExtendedFileRenameInfo, file_name_length) == + offsetof(FILE_RENAME_INFO, FileNameLength)); + static_assert(offsetof(ExtendedFileRenameInfo, file_name) == + offsetof(FILE_RENAME_INFO, FileName)); + + std::error_code path_error; + const std::filesystem::path absolute_destination = + std::filesystem::absolute(destination_path, path_error); + if (path_error) { + ::SetLastError(ERROR_PATH_NOT_FOUND); + return false; + } + const std::wstring &destination = absolute_destination.native(); + if (destination.size() > + (std::numeric_limits::max)() / sizeof(WCHAR)) { + ::SetLastError(ERROR_FILENAME_EXCED_RANGE); + return false; + } + + const size_t destination_bytes = destination.size() * sizeof(WCHAR); + if (destination_bytes > + (std::numeric_limits::max)() - sizeof(ExtendedFileRenameInfo)) { + ::SetLastError(ERROR_FILENAME_EXCED_RANGE); + return false; + } + const size_t rename_info_size = + sizeof(ExtendedFileRenameInfo) + destination_bytes; + + std::unique_ptr rename_buffer( + new (std::nothrow) unsigned char[rename_info_size]); + if (!rename_buffer) { + ::SetLastError(ERROR_NOT_ENOUGH_MEMORY); + return false; + } + std::memset(rename_buffer.get(), 0, rename_info_size); + auto *rename_info = + reinterpret_cast(rename_buffer.get()); + rename_info->flags = kReplaceIfExists | kPosixSemantics; + rename_info->root_directory = nullptr; + rename_info->file_name_length = static_cast(destination_bytes); + std::memcpy(rename_info->file_name, destination.data(), destination_bytes); + + HANDLE source_handle = + ::CreateFileW(source_path.c_str(), DELETE | SYNCHRONIZE, + FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, + nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + if (source_handle == INVALID_HANDLE_VALUE) { + return false; + } + + const BOOL renamed = ::SetFileInformationByHandle( + source_handle, kFileRenameInfoEx, rename_info, + static_cast(rename_info_size)); + const DWORD rename_error = renamed ? ERROR_SUCCESS : ::GetLastError(); + ::CloseHandle(source_handle); + ::SetLastError(rename_error); + return renamed != FALSE; +} +#endif + +SnapshotReplaceResult ReplaceFileAtomically(const std::string &source, + const std::string &destination) { +#if defined(_WIN32) || defined(_WIN64) + const auto source_path = ailego::FileHelper::PathFromUtf8(source); + const auto destination_path = ailego::FileHelper::PathFromUtf8(destination); + // MoveFileExW cannot reliably replace an open destination. DiskAnn keeps + // old snapshots readable until the final in-flight query releases them, so + // use Windows POSIX rename semantics and retain MoveFileExW as a fallback + // for older systems where FileRenameInfoEx is unavailable. + return (ReplaceOpenFileWithPosixSemantics(source_path, destination_path) || + ::MoveFileExW(source_path.c_str(), destination_path.c_str(), + MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH) != + 0) + ? SnapshotReplaceResult::kReplaced + : SnapshotReplaceResult::kNotReplaced; +#else + const size_t separator = destination.rfind('/'); + const std::string parent_directory = + separator == std::string::npos + ? "." + : (separator == 0 ? "/" : destination.substr(0, separator)); + int flags = O_RDONLY; +#ifdef O_DIRECTORY + flags |= O_DIRECTORY; +#endif +#ifdef O_CLOEXEC + flags |= O_CLOEXEC; +#endif + + int directory_fd; + do { + directory_fd = ::open(parent_directory.c_str(), flags); + } while (directory_fd < 0 && errno == EINTR); + if (directory_fd < 0) { + return SnapshotReplaceResult::kNotReplaced; + } + + if (!ailego::File::Rename(source, destination)) { + const int rename_error = errno; + ::close(directory_fd); + errno = rename_error; + return SnapshotReplaceResult::kNotReplaced; + } + + int sync_result; + do { + sync_result = ::fsync(directory_fd); + } while (sync_result != 0 && errno == EINTR); + const int sync_error = sync_result == 0 ? 0 : errno; + ::close(directory_fd); + if (sync_result != 0) { + errno = sync_error; + return SnapshotReplaceResult::kReplacedNotDurable; + } + return SnapshotReplaceResult::kReplaced; +#endif +} + +} // namespace +#endif + +uint32_t DiskAnnIndex::get_doc_count() const { + std::lock_guard lock(mutex_); + if (!is_trained_) { + return static_cast(doc_cache_.size()); + } + return Index::get_doc_count(); +} + #if !DISKANN_SUPPORTED int DiskAnnIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { @@ -44,6 +229,14 @@ int DiskAnnIndex::GenerateHolder() { return core::IndexError_Unsupported; } +int DiskAnnIndex::CommitBuiltSnapshot(bool *snapshot_replaced) { + if (snapshot_replaced != nullptr) { + *snapshot_replaced = false; + } + 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; @@ -93,17 +286,23 @@ int DiskAnnIndex::CreateAndInitStreamer(const BaseIndexParam ¶m) { } param_ = dynamic_cast(param); + if (param_.max_degree <= 0 || param_.list_size <= 0 || + param_.pq_chunk_num < 0) { + LOG_ERROR( + "Invalid DiskAnn parameters: max_degree=%d list_size=%d " + "pq_chunk_num=%d", + param_.max_degree, param_.list_size, param_.pq_chunk_num); + return core::IndexError_InvalidArgument; + } param_.max_degree = std::min(100, param_.max_degree); param_.list_size = std::min(100, param_.list_size); param_.pq_chunk_num = std::min(1024, param_.pq_chunk_num); - proxima_index_params_.set(core::PARAM_DISKANN_BUILDER_MAX_DEGREE, param_.max_degree); proxima_index_params_.set(core::PARAM_DISKANN_BUILDER_LIST_SIZE, param_.list_size); proxima_index_params_.set(core::PARAM_DISKANN_BUILDER_MAX_PQ_CHUNK_NUM, param_.pq_chunk_num); - builder_ = core::IndexFactory::CreateBuilder("DiskAnnBuilder"); streamer_ = core::IndexFactory::CreateStreamer("DiskAnnStreamer"); @@ -187,88 +386,303 @@ int DiskAnnIndex::GenerateHolder() { converter_, &holder_); } +int DiskAnnIndex::CommitBuiltSnapshot(bool *snapshot_replaced) { + if (snapshot_replaced != nullptr) { + *snapshot_replaced = false; + } + if (builder_ == nullptr || file_path_.empty()) { + LOG_ERROR("Cannot commit an uninitialized DiskAnn snapshot"); + return core::IndexError_NoReady; + } + + auto dumper = core::IndexFactory::CreateDumper("FileDumper"); + if (dumper == nullptr) { + LOG_ERROR("Failed to create FileDumper"); + return core::IndexError_Runtime; + } + + const std::string temporary_path = MakeSnapshotTemporaryPath(file_path_); + bool temporary_committed = false; + AILEGO_DEFER([&]() { + if (!temporary_committed && + ailego::FileHelper::IsExist(temporary_path.c_str())) { + ailego::File::Delete(temporary_path); + } + }); + + int ret = dumper->create(temporary_path); + if (ret != 0) { + LOG_ERROR("Failed to create dumper, path: %s, err: %s", + temporary_path.c_str(), core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + ret = builder_->dump(dumper); + if (ret != 0) { + LOG_ERROR("Failed to dump index, path: %s, err: %s", temporary_path.c_str(), + core::IndexError::What(ret)); + dumper->close(); + return core::IndexError_Runtime; + } + ret = dumper->close(); + if (ret != 0) { + LOG_ERROR("Failed to close dumper, path: %s, err: %s", + temporary_path.c_str(), core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + + const core::IndexMeta &build_meta = + converter_ != nullptr ? converter_->meta() : proxima_index_meta_; + auto replacement_storage = + core::IndexFactory::CreateStorage("FileReadStorage"); + auto replacement_streamer = + core::IndexFactory::CreateStreamer("DiskAnnStreamer"); + if (replacement_storage == nullptr || replacement_streamer == nullptr) { + LOG_ERROR("Failed to create replacement DiskAnn reader"); + return core::IndexError_Runtime; + } + + ailego::Params storage_params; + ret = replacement_storage->init(storage_params); + if (ret != 0) { + LOG_ERROR("Failed to initialize replacement storage, err: %s", + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + ret = replacement_streamer->init(build_meta, proxima_index_params_); + if (ret != 0) { + LOG_ERROR("Failed to initialize replacement streamer, err: %s", + core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + ret = replacement_storage->open(temporary_path, false); + if (ret != 0) { + LOG_ERROR("Failed to open replacement storage, path: %s, err: %s", + temporary_path.c_str(), core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + ret = replacement_streamer->open(replacement_storage); + if (ret != 0) { + LOG_ERROR("Failed to validate replacement streamer, path: %s, err: %s", + temporary_path.c_str(), core::IndexError::What(ret)); + return core::IndexError_Runtime; + } + + const SnapshotReplaceResult replace_result = + ReplaceFileAtomically(temporary_path, file_path_); + if (replace_result == SnapshotReplaceResult::kNotReplaced) { + LOG_ERROR("Failed to atomically replace DiskAnn index, path: %s, err: %s", + file_path_.c_str(), + ailego::FileHelper::GetLastErrorString().c_str()); + return core::IndexError_Runtime; + } + temporary_committed = true; + if (snapshot_replaced != nullptr) { + *snapshot_replaced = true; + } + + core::IndexStreamer::Pointer previous_streamer; + core::IndexStorage::Pointer previous_storage; + { + std::lock_guard lock(mutex_); + previous_streamer = exchange_streamer(std::move(replacement_streamer)); + previous_storage = exchange_storage(std::move(replacement_storage)); + } + // Do not explicitly unload the previous streamer. Searches, fetches and + // providers atomically acquire their own shared_ptr snapshot; its file + // reader must remain usable until the last in-flight operation releases it. + // The old streamer and storage clean themselves up when these local owners + // and all reader owners have gone away. + (void)previous_streamer; + (void)previous_storage; + + if (replace_result == SnapshotReplaceResult::kReplacedNotDurable) { + LOG_ERROR( + "DiskAnn snapshot was replaced but its directory entry could not be " + "made durable, path: %s, err: %s", + file_path_.c_str(), ailego::FileHelper::GetLastErrorString().c_str()); + return core::IndexError_WriteData; + } + + return 0; +} + int DiskAnnIndex::add(const VectorData &vector, uint32_t doc_id) { - if (is_trained_) { - LOG_ERROR("this diskann index is trained"); + if (!is_open_) { + LOG_ERROR("Open DiskAnn index before adding vectors"); + return core::IndexError_NoReady; + } + if (is_read_only_) { + LOG_ERROR("Cannot add to a read-only DiskAnn index"); return core::IndexError_Runtime; } if (!std::holds_alternative(vector.vector)) { LOG_ERROR("Invalid vector data"); - return core::IndexError_Runtime; + return core::IndexError_InvalidArgument; } const DenseVector &dense_vector = std::get(vector.vector); - std::string out_vector_buffer = std::string( - static_cast(dense_vector.data), - input_vector_meta_.dimension() * input_vector_meta_.unit_size()); + if (dense_vector.data == nullptr) { + LOG_ERROR("Invalid null vector data"); + return core::IndexError_InvalidArgument; + } + if (doc_id == (std::numeric_limits::max)()) { + LOG_ERROR("Invalid reserved document id: %u", doc_id); + return core::IndexError_OutOfRange; + } - std::lock_guard lock(mutex_); - if (doc_cache_.size() <= doc_id) { - std::string fake_data( - input_vector_meta_.dimension() * input_vector_meta_.unit_size(), 0); - doc_cache_.resize(doc_id + 1, std::make_pair(kInvalidKey, fake_data)); + try { + const size_t vector_size = input_vector_meta_.element_size(); + std::string out_vector_buffer(static_cast(dense_vector.data), + vector_size); + + std::lock_guard lock(mutex_); + if (is_trained_ || is_training_) { + LOG_ERROR("Cannot add vectors while DiskAnn is trained or training"); + return core::IndexError_NoReady; + } + doc_cache_.insert_or_assign(doc_id, std::move(out_vector_buffer)); + } catch (const std::bad_alloc &) { + LOG_ERROR("Not enough memory to cache vector for document id: %u", doc_id); + return core::IndexError_NoMemory; + } catch (const std::length_error &) { + LOG_ERROR("Document id exceeds cache capacity: %u", doc_id); + return core::IndexError_OutOfRange; } - doc_cache_[doc_id] = std::make_pair(doc_id, out_vector_buffer); return 0; } int DiskAnnIndex::train() { + if (!is_open_) { + LOG_ERROR("Open DiskAnn index before training"); + return core::IndexError_NoReady; + } + if (is_read_only_) { + LOG_ERROR("Cannot train a read-only DiskAnn index"); + return core::IndexError_Runtime; + } + { + std::lock_guard lock(mutex_); + if (is_trained_ || is_training_) { + LOG_ERROR("DiskAnn index is already trained or training"); + return core::IndexError_NoReady; + } + is_training_ = true; + } + AILEGO_DEFER([&]() { + std::lock_guard lock(mutex_); + is_training_ = false; + }); + + const core::IndexMeta &build_meta = + converter_ != nullptr ? converter_->meta() : proxima_index_meta_; + auto reset_builder_after_failure = [&]() -> bool { + holder_.reset(); + bool reset_succeeded = true; + if (converter_ != nullptr && converter_->cleanup() != 0) { + LOG_ERROR( + "Failed to release DiskAnn converter result after training " + "failure"); + reset_succeeded = false; + } + if (builder_ == nullptr || builder_->cleanup() != 0 || + builder_->init(build_meta, proxima_index_params_) != 0) { + LOG_ERROR("Failed to reset DiskAnn builder after training failure"); + reset_succeeded = false; + } + return reset_succeeded; + }; + auto return_training_failure = [&](int failure) -> int { + return reset_builder_after_failure() ? failure : core::IndexError_Runtime; + }; + int ret = GenerateHolder(); if (ret != 0) { LOG_ERROR("Failed to generate holder, err: %s", core::IndexError::What(ret)); - return ret; + return return_training_failure(ret); } ret = builder_->train(holder_); if (ret != 0) { LOG_ERROR("Failed to train builder, err: %s", core::IndexError::What(ret)); - return ret; + return return_training_failure(ret); } ret = builder_->build(holder_); if (ret != 0) { LOG_ERROR("Failed to build index, err: %s", core::IndexError::What(ret)); - return ret; - } - auto dumper = core::IndexFactory::CreateDumper("FileDumper"); - if (dumper == nullptr) { - LOG_ERROR("Failed to create FileDumper"); - return core::IndexError_Runtime; + return return_training_failure(ret); } - - ret = dumper->create(file_path_); - if (ret != 0) { - LOG_ERROR("Failed to create dumper, path: %s, err: %s", file_path_.c_str(), - core::IndexError::What(ret)); - return core::IndexError_Runtime; + bool snapshot_replaced = false; + ret = CommitBuiltSnapshot(&snapshot_replaced); + if (ret != 0 && !snapshot_replaced) { + return return_training_failure(ret); } - ret = builder_->dump(dumper); - if (ret != 0) { - LOG_ERROR("Failed to dump index, path: %s, err: %s", file_path_.c_str(), - core::IndexError::What(ret)); - return core::IndexError_Runtime; + // The committed streamer owns the searchable snapshot. Drop all build-time + // copies immediately so a trained mobile index does not retain the input + // cache, holder, and builder entity for the rest of its lifetime. + holder_.reset(); + { + std::lock_guard lock(mutex_); + is_trained_ = true; + decltype(doc_cache_) empty_cache; + doc_cache_.swap(empty_cache); } - dumper->close(); - ret = storage_->open(file_path_, false); - if (ret != 0) { - LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), - core::IndexError::What(ret)); - return core::IndexError_Runtime; + if (builder_->cleanup() != 0) { + LOG_WARN("Failed to release DiskAnn builder memory after training"); } - if (streamer_ == nullptr || streamer_->open(storage_) != 0) { - LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); - return core::IndexError_Runtime; + if (converter_ != nullptr && converter_->cleanup() != 0) { + LOG_WARN("Failed to release DiskAnn converter memory after training"); } - is_trained_ = true; - return 0; + return ret; } int DiskAnnIndex::_dense_fetch(const uint32_t doc_id, VectorDataBuffer *vector_data_buffer) { if (is_trained_) { - return Index::_dense_fetch(doc_id, vector_data_buffer); + const auto streamer = streamer_snapshot(); + if (streamer == nullptr) { + return core::IndexError_NoReady; + } + auto &context = acquire_context(); + if (context == nullptr) { + LOG_ERROR("Failed to acquire DiskAnn fetch context"); + return core::IndexError_Runtime; + } + + std::string stored_vector; + const int ret = streamer->get_vector(doc_id, context, stored_vector); + context->reset(); + if (ret != 0) { + return ret; + } + const size_t expected_vector_size = streamer_vector_meta_.element_size(); + if (stored_vector.size() != expected_vector_size) { + LOG_ERROR("Invalid fetched vector size: %zu, expected: %zu", + stored_vector.size(), expected_vector_size); + return core::IndexError_InvalidFormat; + } + + DenseVectorBuffer dense_vector_buffer; + if (reformer_ != nullptr) { + dense_vector_buffer.data.resize(input_vector_meta_.element_size()); + if (reformer_->revert(stored_vector.data(), streamer_vector_meta_, + &dense_vector_buffer.data) != 0) { + LOG_ERROR("Failed to revert fetched DiskAnn vector"); + return core::IndexError_Runtime; + } + } else { + dense_vector_buffer.data = std::move(stored_vector); + } + vector_data_buffer->vector_buffer = std::move(dense_vector_buffer); + return 0; } else { + std::lock_guard lock(mutex_); + const auto iter = doc_cache_.find(doc_id); + if (iter == doc_cache_.end()) { + LOG_ERROR("Vector id does not exist: %u", doc_id); + return core::IndexError_NoExist; + } DenseVectorBuffer dense_vector_buffer; std::string &out_vector_buffer = dense_vector_buffer.data; - out_vector_buffer = doc_cache_[doc_id].second; + out_vector_buffer = iter->second; vector_data_buffer->vector_buffer = std::move(dense_vector_buffer); return 0; } @@ -309,7 +723,12 @@ int DiskAnnIndex::_prepare_for_search( params.set( core::PARAM_DISKANN_SEARCHER_LIST_SIZE, std::max(diskann_search_param->topk, diskann_search_param->list_size)); - context->update(params); + const int ret = context->update(params); + if (ret != 0) { + LOG_ERROR("Failed to update DiskAnn search context: %s", + core::IndexError::What(ret)); + return ret; + } return 0; } @@ -317,34 +736,90 @@ int DiskAnnIndex::_prepare_for_search( int DiskAnnIndex::merge(const std::vector &indexes, const IndexFilter &filter, const MergeOptions &options) { - int pre_ret = Index::merge(indexes, filter, options); - if (pre_ret != 0) { - return pre_ret; + if (indexes.empty()) { + return core::IndexError_Success; } - auto dumper = core::IndexFactory::CreateDumper("FileDumper"); - - dumper->create(file_path_); - int ret = builder_->dump(dumper); - if (ret != 0) { - LOG_ERROR("Failed to dump index, path: %s, err: %s", file_path_.c_str(), - core::IndexError::What(ret)); + if (!is_open_) { + LOG_ERROR("Open DiskAnn index before merging"); + return core::IndexError_NoReady; + } + if (is_read_only_) { + LOG_ERROR("Cannot merge into a read-only DiskAnn index"); return core::IndexError_Runtime; } - dumper->close(); + bool was_trained = false; + { + std::lock_guard lock(mutex_); + if (is_training_) { + LOG_ERROR("DiskAnn index is already training or merging"); + return core::IndexError_NoReady; + } + is_training_ = true; + was_trained = is_trained_; + } + AILEGO_DEFER([&]() { + std::lock_guard lock(mutex_); + is_training_ = false; + }); + + const core::IndexMeta &build_meta = + converter_ != nullptr ? converter_->meta() : proxima_index_meta_; + auto rollback_training_state = [&]() { + { + std::lock_guard lock(mutex_); + is_trained_ = was_trained; + } + holder_.reset(); + if (converter_ != nullptr && converter_->cleanup() != 0) { + LOG_ERROR( + "Failed to release DiskAnn converter result after merge " + "failure"); + } + if (builder_ == nullptr || builder_->cleanup() != 0 || + builder_->init(build_meta, proxima_index_params_) != 0) { + LOG_ERROR("Failed to reset DiskAnn builder after merge failure"); + } + }; + + // A DiskAnn builder is single-use. Reinitialize it before rebuilding an + // already trained target so repeated merge calls behave like other indexes. + if (was_trained) { + if (builder_ == nullptr || builder_->cleanup() != 0) { + LOG_ERROR("Failed to reset DiskAnn builder before merge"); + return core::IndexError_Runtime; + } + if (builder_->init(build_meta, proxima_index_params_) != 0) { + LOG_ERROR("Failed to reinitialize DiskAnn builder before merge"); + return core::IndexError_Runtime; + } + } - ret = storage_->open(file_path_, false); - if (ret != 0) { - LOG_ERROR("Failed to open storage, path: %s, err: %s", file_path_.c_str(), - core::IndexError::What(ret)); - return core::IndexError_Runtime; + int pre_ret = Index::merge(indexes, filter, options); + if (pre_ret != 0) { + rollback_training_state(); + return pre_ret; } - if (streamer_ == nullptr || streamer_->open(storage_) != 0) { - LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str()); - return core::IndexError_Runtime; + bool snapshot_replaced = false; + const int ret = CommitBuiltSnapshot(&snapshot_replaced); + if (ret != 0 && !snapshot_replaced) { + rollback_training_state(); + return ret; } - is_trained_ = true; - return 0; + holder_.reset(); + { + std::lock_guard lock(mutex_); + is_trained_ = true; + decltype(doc_cache_) empty_cache; + doc_cache_.swap(empty_cache); + } + if (builder_->cleanup() != 0) { + LOG_WARN("Failed to release DiskAnn builder memory after merge"); + } + if (converter_ != nullptr && converter_->cleanup() != 0) { + LOG_WARN("Failed to release DiskAnn converter memory after merge"); + } + return ret; } #endif // DISKANN_SUPPORTED diff --git a/src/core/interface/indexes/holder_builder.h b/src/core/interface/indexes/holder_builder.h index 83021adbc..5fbc96c70 100644 --- a/src/core/interface/indexes/holder_builder.h +++ b/src/core/interface/indexes/holder_builder.h @@ -30,11 +30,10 @@ namespace zvec::core_interface { inline constexpr uint64_t kInvalidKey = std::numeric_limits::max(); -template -inline int BuildMultiPassHolderImpl( - uint32_t dimension, - const std::vector> &doc_cache, - core::IndexHolder::Pointer *holder_out) { +template +inline int BuildMultiPassHolderImpl(uint32_t dimension, + const DocCache &doc_cache, + core::IndexHolder::Pointer *holder_out) { auto holder = std::make_shared>(dimension); for (const auto &doc : doc_cache) { @@ -51,11 +50,11 @@ inline int BuildMultiPassHolderImpl( return 0; } -inline int BuildMultiPassHolder( - DataType data_type, uint32_t dimension, - const std::vector> &doc_cache, - const core::IndexConverter::Pointer &converter, - core::IndexHolder::Pointer *holder) { +template +inline int BuildMultiPassHolder(DataType data_type, uint32_t dimension, + const DocCache &doc_cache, + const core::IndexConverter::Pointer &converter, + core::IndexHolder::Pointer *holder) { int ret = 0; switch (data_type) { case DataType::DT_FP16: @@ -78,8 +77,17 @@ inline int BuildMultiPassHolder( return ret; } if (converter) { - core::IndexConverter::TrainAndTransform(converter, *holder); - *holder = converter->result(); + ret = core::IndexConverter::TrainAndTransform(converter, *holder); + if (ret != 0) { + LOG_ERROR("Failed to train and transform holder, ret=%d", ret); + return ret; + } + auto converted_holder = converter->result(); + if (!converted_holder) { + LOG_ERROR("Converter returned no result holder"); + return core::IndexError_Runtime; + } + *holder = std::move(converted_holder); } return 0; } diff --git a/src/core/mixed_reducer/mixed_streamer_reducer.cc b/src/core/mixed_reducer/mixed_streamer_reducer.cc index 01d18ff2f..38c9767ef 100644 --- a/src/core/mixed_reducer/mixed_streamer_reducer.cc +++ b/src/core/mixed_reducer/mixed_streamer_reducer.cc @@ -26,6 +26,45 @@ namespace zvec { namespace core { +namespace { + +bool matches_query_meta(const IndexMeta &meta, + const IndexQueryMeta &query_meta) { + if (meta.meta_type() != query_meta.meta_type() || + meta.data_type() != query_meta.data_type() || + meta.unit_size() != query_meta.unit_size()) { + return false; + } + // Sparse records carry their logical length per document. Their metadata + // dimension and element size are not the size of an individual record. + return meta.meta_type() == IndexMeta::MetaType::MT_SPARSE || + (meta.dimension() == query_meta.dimension() && + meta.element_size() == query_meta.element_size()); +} + +bool matches_index_meta(const IndexQueryMeta &query_meta, + const IndexMeta &meta) { + return matches_query_meta(meta, query_meta); +} + +bool checked_add(uint64_t lhs, uint64_t rhs, uint64_t *result) { + if (rhs > (std::numeric_limits::max)() - lhs) { + return false; + } + *result = lhs + rhs; + return true; +} + +bool checked_multiply(size_t lhs, size_t rhs, size_t *result) { + if (lhs != 0 && rhs > (std::numeric_limits::max)() / lhs) { + return false; + } + *result = lhs * rhs; + return true; +} + +} // namespace + int MixedStreamerReducer::init(const ailego::Params ¶ms) { enable_pk_rewrite_ = params.get_as_bool(PARAM_MIXED_STREAMER_REDUCER_ENABLE_PK_REWRITE); @@ -44,15 +83,31 @@ int MixedStreamerReducer::init(const ailego::Params ¶ms) { } int MixedStreamerReducer::cleanup(void) { + int ret = 0; streamers_.clear(); - target_streamer_->cleanup(); - - target_builder_->cleanup(); + source_streamers_reformers_.clear(); + if (target_streamer_ != nullptr) { + ret = target_streamer_->cleanup(); + } + if (target_builder_ != nullptr) { + const int builder_ret = target_builder_->cleanup(); + if (ret == 0) { + ret = builder_ret; + } + } + target_streamer_.reset(); + target_streamer_reformer_.reset(); + target_builder_.reset(); + target_builder_converter_.reset(); doc_cache_.clear(); + mt_list_.reset(); + mt_list_.resume_consume(); + sparse_mt_list_.reset(); + sparse_mt_list_.resume_consume(); stats_.clear_attributes(); state_ = STATE_UNINITED; - return 0; + return ret; } int MixedStreamerReducer::set_target_streamer_wiht_info( @@ -64,6 +119,11 @@ int MixedStreamerReducer::set_target_streamer_wiht_info( LOG_ERROR("Set target streamer after init"); return IndexError_Uninitialized; } + if (!streamer || + original_query_meta.meta_type() != streamer->meta().meta_type()) { + LOG_ERROR("Invalid target streamer or original query metadata"); + return IndexError_InvalidArgument; + } target_builder_ = builder; target_streamer_ = streamer; @@ -90,35 +150,30 @@ int MixedStreamerReducer::feed_streamer_with_reformer( return IndexError_InvalidArgument; } - auto check_datatype = [&](const IndexMeta & /*target_meta*/, - const IndexMeta &source_meta) -> bool { - if (!streamers_.empty()) { - auto &last_meta = streamers_.back()->meta(); - return last_meta.data_type() == source_meta.data_type() && - last_meta.dimension() == source_meta.dimension() && - last_meta.unit_size() == source_meta.unit_size(); - } - // TODO: check target meta - return true; - }; - - auto check_other = [&](const IndexMeta &target_meta, - const IndexMeta &source_meta) -> bool { - return target_meta.meta_type() == source_meta.meta_type(); - // when create a new index, there is a case that ip_flat merged into l2_hnsw - // target_meta.metric_name() == source_meta.metric_name(); - }; - - if (!(check_datatype(target_streamer_->meta(), streamer->meta()) && - check_other(target_streamer_->meta(), streamer->meta()))) { + const IndexMeta &target_meta = target_streamer_->meta(); + const IndexMeta &source_meta = streamer->meta(); + if (target_meta.meta_type() != source_meta.meta_type()) { LOG_ERROR("Streamer meta mismatch"); return IndexError_InvalidArgument; } - if (streamers_.empty()) { - is_target_and_source_same_reformer_ = - target_streamer_->meta().reformer_name() == - streamer->meta().reformer_name(); + const bool source_is_encoded = !source_meta.reformer_name().empty(); + const bool target_is_encoded = !target_meta.reformer_name().empty(); + const bool source_is_decodable = + source_is_encoded ? reformer != nullptr + : matches_query_meta(source_meta, original_query_meta_); + bool compatible = source_is_decodable; + // Builders consume the original representation and retrain their target + // converter after all source records have been decoded. + if (target_builder_ == nullptr) { + compatible = compatible && + (target_is_encoded + ? target_streamer_reformer_ != nullptr + : matches_query_meta(target_meta, original_query_meta_)); + } + if (!compatible) { + LOG_ERROR("Streamer vector representation mismatch"); + return IndexError_InvalidArgument; } streamers_.push_back(streamer); @@ -148,15 +203,51 @@ int MixedStreamerReducer::reduce(const IndexFilter &filter) { // TODO: use id instead of key // When merging into a non-empty target (e.g. reusing one input as base), // append new docs after the existing ones instead of overwriting from 0. - uint32_t id_offset = 0; - uint32_t next_id = 0; + uint64_t id_offset = 0; + uint64_t next_id = 0; + auto find_next_target_id = [&next_id](const auto &provider) -> int { + const uint64_t target_count = provider->count(); + next_id = (std::max)(next_id, target_count); + if (target_count == 0) { + return 0; + } + auto iterator = provider->create_iterator(); + if (!iterator) { + LOG_ERROR("Failed to create target provider iterator"); + return IndexError_Runtime; + } + while (iterator->is_valid()) { + const uint64_t key = iterator->key(); + if (key == (std::numeric_limits::max)()) { + LOG_ERROR("Invalid target vector key"); + return IndexError_InvalidFormat; + } + next_id = (std::max)(next_id, key + 1); + iterator->next(); + } + return 0; + }; if (target_builder_ == nullptr) { if (is_sparse_) { auto provider = target_streamer_->create_sparse_provider(); - if (provider) next_id = provider->count(); + if (!provider) { + LOG_ERROR("Failed to create target sparse provider"); + return IndexError_Runtime; + } + const int ret = find_next_target_id(provider); + if (ret != 0) { + return ret; + } } else { auto provider = target_streamer_->create_provider(); - if (provider) next_id = provider->count(); + if (!provider) { + LOG_ERROR("Failed to create target provider"); + return IndexError_Runtime; + } + const int ret = find_next_target_id(provider); + if (ret != 0) { + return ret; + } } } @@ -168,8 +259,23 @@ int MixedStreamerReducer::reduce(const IndexFilter &filter) { for (size_t i = 0; i < streamers_.size(); i++) { // due to filter, producing can't be parallel - read_results[i] = read_sparse_vec(i, filter, id_offset, &next_id); - id_offset += streamers_[i]->create_sparse_provider()->count(); + auto provider = streamers_[i]->create_sparse_provider(); + if (!provider) { + LOG_ERROR("Failed to create source sparse provider, index=%zu", i); + read_results[i] = IndexError_Runtime; + break; + } + uint64_t source_span = 0; + read_results[i] = read_sparse_vec(i, provider, filter, id_offset, + &next_id, &source_span); + if (read_results[i] != 0) { + break; + } + if (!checked_add(id_offset, source_span, &id_offset)) { + LOG_ERROR("Source sparse vector key range overflows"); + read_results[i] = IndexError_InvalidFormat; + break; + } } sparse_mt_list_.done(); @@ -181,8 +287,23 @@ int MixedStreamerReducer::reduce(const IndexFilter &filter) { } for (size_t i = 0; i < streamers_.size(); i++) { - read_results[i] = read_vec(i, filter, id_offset, &next_id); - id_offset += streamers_[i]->create_provider()->count(); + auto provider = streamers_[i]->create_provider(); + if (!provider) { + LOG_ERROR("Failed to create source provider, index=%zu", i); + read_results[i] = IndexError_Runtime; + break; + } + uint64_t source_span = 0; + read_results[i] = + read_vec(i, provider, filter, id_offset, &next_id, &source_span); + if (read_results[i] != 0) { + break; + } + if (!checked_add(id_offset, source_span, &id_offset)) { + LOG_ERROR("Source vector key range overflows"); + read_results[i] = IndexError_InvalidFormat; + break; + } } mt_list_.done(); @@ -204,8 +325,6 @@ int MixedStreamerReducer::reduce(const IndexFilter &filter) { return IndexError_Runtime; } - stats_.set_reduced_costtime(timer.seconds()); - state_ = STATE_REDUCE; if (target_builder_ != nullptr) { int ret = IndexBuild(); if (ret != 0) { @@ -214,6 +333,8 @@ int MixedStreamerReducer::reduce(const IndexFilter &filter) { } } + stats_.set_reduced_costtime(timer.seconds()); + state_ = STATE_REDUCE; LOG_INFO("End brute force reduce. cost time: [%zu]s", (size_t)timer.seconds()); return 0; @@ -227,12 +348,17 @@ int MixedStreamerReducer::dump(const IndexDumper::Pointer &dumper) { return IndexError_NoReady; } + if (!dumper) { + LOG_ERROR("Dumper is null"); + return IndexError_InvalidArgument; + } + ailego::ElapsedTime timer; int ret = 0; if (target_builder_ != nullptr) { - target_builder_->dump(dumper); + ret = target_builder_->dump(dumper); } else { - target_streamer_->dump(dumper); + ret = target_streamer_->dump(dumper); } if (ret == IndexError_NotImplemented) { LOG_WARN("Dump index not implemented"); @@ -244,52 +370,116 @@ int MixedStreamerReducer::dump(const IndexDumper::Pointer &dumper) { } int MixedStreamerReducer::read_vec(size_t source_streamer_index, + const IndexProvider::Pointer &provider, const IndexFilter &filter, - const uint32_t id_offset, - uint32_t *next_id) { + uint64_t id_offset, uint64_t *next_id, + uint64_t *source_span) { const auto &streamer = streamers_[source_streamer_index]; const auto &reformer = source_streamers_reformers_[source_streamer_index]; const IndexQueryMeta source_streamer_query_meta{streamer->meta().data_type(), streamer->meta().dimension()}; - - bool need_revert = (target_streamer_->meta().reformer_name() != - streamer->meta().reformer_name() && - reformer != nullptr); - if (target_builder_ && reformer) { - need_revert = true; + // A reformer name identifies an algorithm, not its trained state. Separate + // indexes can use the same name with different scale/bias values or rotation + // matrices. Always return source records to the original representation and + // then encode them with the target reformer instead of copying encoded bytes. + const bool need_revert = !streamer->meta().reformer_name().empty(); + const bool need_convert = target_builder_ == nullptr && + !target_streamer_->meta().reformer_name().empty(); + + if (!provider) { + LOG_ERROR("Source provider is null, index=%zu", source_streamer_index); + return IndexError_Runtime; } - - IndexProvider::Pointer provider = streamer->create_provider(); + *source_span = provider->count(); IndexProvider::Iterator::Pointer iterator = provider->create_iterator(); + if (!iterator) { + LOG_ERROR("Failed to create source provider iterator, index=%zu", + source_streamer_index); + return IndexError_Runtime; + } while (iterator->is_valid()) { if (stop_flag_ != nullptr && stop_flag_->load(std::memory_order_relaxed)) { LOG_DEBUG("read_vec cancelled."); return 0; } - if (filter(iterator->key() + (uint64_t)id_offset)) { + const uint64_t source_key = iterator->key(); + if (source_key == (std::numeric_limits::max)()) { + LOG_ERROR("Invalid source vector key"); + return IndexError_InvalidFormat; + } + *source_span = (std::max)(*source_span, source_key + 1); + uint64_t global_id = 0; + if (!checked_add(id_offset, source_key, &global_id)) { + LOG_ERROR("Source vector key overflows global id range"); + return IndexError_InvalidFormat; + } + if (filter(global_id)) { (*stats_.mutable_filtered_count())++; iterator->next(); continue; } - std::vector bytes; + const void *vector_data = iterator->data(); + if (!vector_data) { + LOG_ERROR("Failed to read source vector, index=%zu key=%zu", + source_streamer_index, static_cast(iterator->key())); + return IndexError_ReadData; + } + std::string reverted_vector; if (need_revert) { - std::string new_vector; - if (reformer->revert(iterator->data(), source_streamer_query_meta, - &new_vector) != 0) { + if (reformer->revert(vector_data, source_streamer_query_meta, + &reverted_vector) != 0) { LOG_ERROR("Failed to revert the vector"); return IndexError_Runtime; } - bytes.resize(new_vector.size()); - memcpy(bytes.data(), new_vector.data(), bytes.size()); - } else { - // TODO: eliminate the copy - bytes.resize(provider->element_size()); - memcpy(bytes.data(), iterator->data(), bytes.size()); + if (reverted_vector.size() != original_query_meta_.element_size()) { + LOG_ERROR("Reverted vector has an invalid size: actual=%zu expected=%u", + reverted_vector.size(), original_query_meta_.element_size()); + return IndexError_Mismatch; + } + vector_data = reverted_vector.data(); + } + + std::string converted_vector; + if (need_convert) { + IndexQueryMeta converted_meta; + if (target_streamer_reformer_->convert(vector_data, original_query_meta_, + &converted_vector, + &converted_meta) != 0) { + LOG_ERROR("Failed to convert vector into target representation"); + return IndexError_Runtime; + } + if (!matches_index_meta(converted_meta, target_streamer_->meta())) { + LOG_ERROR("Converted vector metadata does not match target streamer"); + return IndexError_Mismatch; + } + vector_data = converted_vector.data(); + } + + const size_t expected_size = target_builder_ != nullptr + ? original_query_meta_.element_size() + : target_streamer_->meta().element_size(); + const size_t source_size = + need_convert + ? converted_vector.size() + : (need_revert ? reverted_vector.size() : provider->element_size()); + if (source_size != expected_size) { + LOG_ERROR("Source vector has an invalid size: actual=%zu expected=%zu", + source_size, expected_size); + return IndexError_Mismatch; + } + + std::vector bytes(expected_size); + if (!bytes.empty()) { + memcpy(bytes.data(), vector_data, bytes.size()); } // TODO: use id instead of key + if (*next_id > (std::numeric_limits::max)()) { + LOG_ERROR("Target vector id overflows uint32 range"); + return IndexError_InvalidFormat; + } if (!mt_list_.produce(VectorItem((*next_id)++, std::move(bytes)))) { LOG_ERROR("Produce vector to queue failed. key[%lu]", (size_t)iterator->key()); @@ -310,8 +500,6 @@ void MixedStreamerReducer::add_vec(int *result) { auto target_streamer_query_meta = IndexQueryMeta{ IndexMeta::MetaType::MT_DENSE, target_streamer_->meta().data_type(), target_streamer_->meta().dimension()}; - const bool need_convert = (!is_target_and_source_same_reformer_) && - target_streamer_reformer_ != nullptr; AILEGO_DEFER([&]() { // make producer quit @@ -325,31 +513,10 @@ void MixedStreamerReducer::add_vec(int *result) { return; } - const void *vector = vector_item.vec_.data(); - std::string new_vector; - - - if (need_convert) { - IndexQueryMeta new_meta; - if (target_streamer_reformer_->convert(vector, original_query_meta_, - &new_vector, &new_meta) != 0) { - LOG_ERROR("Failed to transform vector"); - *result = IndexError_Runtime; - return; - } - vector = new_vector.data(); - } - // 1. no reformer: target_streamer_query_meta_ = original_query_meta_ - // 2. has reformer, matched(need_convert = false): use - // target_streamer_query_meta_ - // 3. has reformer, not matched(need_convert = true): use - // target_streamer_query_meta_ - - // TODO: use id instead of key int ret = target_streamer_->add_with_id_impl( - (uint32_t)vector_item.pkey_, vector, target_streamer_query_meta, - target_streamer_context); + (uint32_t)vector_item.pkey_, vector_item.vec_.data(), + target_streamer_query_meta, target_streamer_context); if (ret != 0) { LOG_ERROR("Insert target streamer failed. ret[%d] reason[%s] pkey[%zu]", ret, IndexError::What(ret), (size_t)vector_item.pkey_); @@ -399,9 +566,6 @@ void MixedStreamerReducer::add_sparse_vec(int *result) { target_streamer_->meta().data_type(), }; - auto need_convert = !is_target_and_source_same_reformer_ && - target_streamer_reformer_ != nullptr; - AILEGO_DEFER([&]() { // make producer quit sparse_mt_list_.done(); @@ -417,20 +581,6 @@ void MixedStreamerReducer::add_sparse_vec(int *result) { auto indices = sparse_vector_item.sparse_indices_.data(); auto values = sparse_vector_item.sparse_values_.data(); - std::string converted_sparse_values_buffer; - if (need_convert) { - IndexQueryMeta new_meta; - if (target_streamer_reformer_->convert( - sparse_count, indices, values, original_query_meta_, - &converted_sparse_values_buffer, &new_meta) != 0) { - LOG_ERROR("Failed to transform vector"); - *result = IndexError_Runtime; - return; - } - values = converted_sparse_values_buffer.data(); - target_streamer_query_meta = new_meta; - } - // TODO: use id instead of key int ret = target_streamer_->add_with_id_impl( (uint32_t)sparse_vector_item.pkey_, sparse_count, indices, values, @@ -449,59 +599,137 @@ void MixedStreamerReducer::add_sparse_vec(int *result) { } -int MixedStreamerReducer::read_sparse_vec(size_t source_streamer_index, - const IndexFilter &filter, - const uint32_t id_offset, - uint32_t *next_id) { +int MixedStreamerReducer::read_sparse_vec( + size_t source_streamer_index, + const IndexStreamer::SparseProvider::Pointer &provider, + const IndexFilter &filter, uint64_t id_offset, uint64_t *next_id, + uint64_t *source_span) { const auto &streamer = streamers_[source_streamer_index]; const auto &reformer = source_streamers_reformers_[source_streamer_index]; - const bool need_revert = - !is_target_and_source_same_reformer_ && reformer != nullptr; + const bool need_revert = !streamer->meta().reformer_name().empty(); + const bool need_convert = target_builder_ == nullptr && + !target_streamer_->meta().reformer_name().empty(); - IndexStreamer::SparseProvider::Pointer provider = - streamer->create_sparse_provider(); + if (!provider) { + LOG_ERROR("Source sparse provider is null, index=%zu", + source_streamer_index); + return IndexError_Runtime; + } + *source_span = provider->count(); IndexStreamer::SparseProvider::Iterator::Pointer iterator = provider->create_iterator(); + if (!iterator) { + LOG_ERROR("Failed to create source sparse provider iterator, index=%zu", + source_streamer_index); + return IndexError_Runtime; + } while (iterator->is_valid()) { if (stop_flag_ != nullptr && stop_flag_->load(std::memory_order_relaxed)) { LOG_DEBUG("read_sparse_vec cancelled."); return 0; } - if (filter(iterator->key() + (uint64_t)id_offset)) { + const uint64_t source_key = iterator->key(); + if (source_key == (std::numeric_limits::max)()) { + LOG_ERROR("Invalid source sparse vector key"); + return IndexError_InvalidFormat; + } + *source_span = (std::max)(*source_span, source_key + 1); + uint64_t global_id = 0; + if (!checked_add(id_offset, source_key, &global_id)) { + LOG_ERROR("Source sparse vector key overflows global id range"); + return IndexError_InvalidFormat; + } + if (filter(global_id)) { (*stats_.mutable_filtered_count())++; iterator->next(); continue; } - auto sparse_count = iterator->sparse_count(); + const auto sparse_count = iterator->sparse_count(); + const uint32_t *const source_indices = iterator->sparse_indices(); + const void *const source_values = iterator->sparse_data(); + if (sparse_count > 0 && + (source_indices == nullptr || source_values == nullptr)) { + LOG_ERROR("Failed to read source sparse vector, index=%zu", + source_streamer_index); + return IndexError_ReadData; + } std::vector sparse_indices(sparse_count); + if (!sparse_indices.empty()) { + memcpy(sparse_indices.data(), source_indices, + sparse_indices.size() * sizeof(uint32_t)); + } + const void *values = source_values; std::string sparse_values; if (need_revert) { - std::string new_sparse_values; - if (reformer->revert(iterator->sparse_count(), iterator->sparse_indices(), - iterator->sparse_data(), + if (reformer->revert(sparse_count, source_indices, source_values, { IndexMeta::MetaType::MT_SPARSE, streamer->meta().data_type(), }, - &new_sparse_values) != 0) { + &sparse_values) != 0) { LOG_ERROR("Failed to revert the sparse vector"); return IndexError_Runtime; } - sparse_values = std::move(new_sparse_values); - } else { - sparse_values.resize(sparse_count * streamer->meta().unit_size()); - memcpy(sparse_values.data(), iterator->sparse_data(), - sparse_values.size()); + values = sparse_values.data(); } - // TODO: eliminate the copy - memcpy(sparse_indices.data(), iterator->sparse_indices(), - sparse_indices.size() * sizeof(uint32_t)); + std::string converted_sparse_values; + if (need_convert) { + IndexQueryMeta converted_meta; + if (target_streamer_reformer_->convert( + sparse_count, source_indices, values, original_query_meta_, + &converted_sparse_values, &converted_meta) != 0) { + LOG_ERROR("Failed to convert sparse vector into target representation"); + return IndexError_Runtime; + } + if (!matches_index_meta(converted_meta, target_streamer_->meta())) { + LOG_ERROR( + "Converted sparse vector metadata does not match target streamer"); + return IndexError_Mismatch; + } + values = converted_sparse_values.data(); + } + + const size_t expected_unit_size = + target_builder_ != nullptr ? original_query_meta_.unit_size() + : target_streamer_->meta().unit_size(); + size_t expected_size = 0; + if (!checked_multiply(sparse_count, expected_unit_size, &expected_size)) { + LOG_ERROR("Sparse vector size overflows"); + return IndexError_InvalidFormat; + } + size_t raw_source_size = 0; + if (!checked_multiply(sparse_count, streamer->meta().unit_size(), + &raw_source_size)) { + LOG_ERROR("Source sparse vector size overflows"); + return IndexError_InvalidFormat; + } + const size_t source_size = + need_convert ? converted_sparse_values.size() + : (need_revert ? sparse_values.size() : raw_source_size); + if (source_size != expected_size) { + LOG_ERROR( + "Source sparse vector has an invalid size: actual=%zu expected=%zu", + source_size, expected_size); + return IndexError_Mismatch; + } + if (!need_revert && !need_convert) { + sparse_values.resize(expected_size); + if (!sparse_values.empty()) { + memcpy(sparse_values.data(), values, sparse_values.size()); + } + } else if (need_convert) { + sparse_values = std::move(converted_sparse_values); + } // TODO: use id instead of key + if (*next_id > (std::numeric_limits::max)()) { + LOG_ERROR("Target sparse vector id overflows uint32 range"); + return IndexError_InvalidFormat; + } if (!sparse_mt_list_.produce(SparseVectorItem((*next_id)++, std::move(sparse_indices), std::move(sparse_values)))) { @@ -578,9 +806,17 @@ int MixedStreamerReducer::IndexBuild() { return core::IndexError_Runtime; } if (target_builder_converter_) { - core::IndexConverter::TrainAndTransform(target_builder_converter_, - target_holder); + int ret = core::IndexConverter::TrainAndTransform(target_builder_converter_, + target_holder); + if (ret != 0) { + LOG_ERROR("Failed to convert target holder, ret=%d", ret); + return ret; + } target_holder = target_builder_converter_->result(); + if (!target_holder) { + LOG_ERROR("Target converter returned no result holder"); + return core::IndexError_Runtime; + } } int ret = target_builder_->train(target_holder); if (ret != 0) { diff --git a/src/core/mixed_reducer/mixed_streamer_reducer.h b/src/core/mixed_reducer/mixed_streamer_reducer.h index ec4c62406..f1147c254 100644 --- a/src/core/mixed_reducer/mixed_streamer_reducer.h +++ b/src/core/mixed_reducer/mixed_streamer_reducer.h @@ -57,12 +57,16 @@ class MixedStreamerReducer : public IndexStreamerReducer { const IndexReformer::Pointer reformer) override; private: - int read_vec(size_t source_streamer_index, const IndexFilter &filter, - const uint32_t id_offset, uint32_t *next_id); + int read_vec(size_t source_streamer_index, + const IndexProvider::Pointer &provider, + const IndexFilter &filter, uint64_t id_offset, uint64_t *next_id, + uint64_t *source_span); void add_vec(int *result); void add_vec_with_builder(int *result); - int read_sparse_vec(size_t source_streamer_index, const IndexFilter &filter, - const uint32_t id_offset, uint32_t *next_id); + int read_sparse_vec(size_t source_streamer_index, + const IndexStreamer::SparseProvider::Pointer &provider, + const IndexFilter &filter, uint64_t id_offset, + uint64_t *next_id, uint64_t *source_span); void add_sparse_vec(int *result); void PushToDocCache(const IndexQueryMeta &meta, uint32_t doc_id, @@ -97,7 +101,6 @@ class MixedStreamerReducer : public IndexStreamerReducer { ailego::Params params_; IndexStreamer::Pointer target_streamer_{nullptr}; IndexReformer::Pointer target_streamer_reformer_{nullptr}; - bool is_target_and_source_same_reformer_{false}; IndexQueryMeta original_query_meta_{}; std::vector streamers_; diff --git a/src/core/quantizer/binary_converter.cc b/src/core/quantizer/binary_converter.cc index 7ef5672fb..a33039879 100644 --- a/src/core/quantizer/binary_converter.cc +++ b/src/core/quantizer/binary_converter.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. #include +#include #include #include #include @@ -29,14 +30,13 @@ class BinaryConverterHolder : public IndexHolder { class Iterator : public IndexHolder::Iterator { public: //! Constructor - Iterator(const BinaryConverterHolder *owner, + Iterator(size_t dimension, + const std::shared_ptr &quantizer, IndexHolder::Iterator::Pointer &&iter) - : buffer_(ailego::BinaryQuantizer::EncodedSizeInBinary32( - owner->dimension()), - 0), + : buffer_(ailego::BinaryQuantizer::EncodedSizeInBinary32(dimension), 0), front_iter_(std::move(iter)), - quantizer_(owner->quantizer_), - dim_{owner->dimension()} { + quantizer_(quantizer), + dim_{dimension} { this->encode_record(); } @@ -45,7 +45,7 @@ class BinaryConverterHolder : public IndexHolder { //! Retrieve pointer of data const void *data(void) const override { - return buffer_.data(); + return data_valid_ ? buffer_.data() : nullptr; } //! Test if the iterator is valid @@ -67,10 +67,16 @@ class BinaryConverterHolder : public IndexHolder { private: //! Encode the data by quantizer inline void encode_record(void) { - if (front_iter_->is_valid()) { - const float *vec = reinterpret_cast(front_iter_->data()); - quantizer_->encode(vec, dim_ / 2, buffer_.data()); + data_valid_ = false; + if (!front_iter_->is_valid()) { + return; } + const float *vec = static_cast(front_iter_->data()); + if (!vec || !quantizer_) { + return; + } + quantizer_->encode(vec, dim_, buffer_.data()); + data_valid_ = true; } //! Members @@ -78,6 +84,7 @@ class BinaryConverterHolder : public IndexHolder { IndexHolder::Iterator::Pointer front_iter_{}; std::shared_ptr quantizer_{}; size_t dim_{0u}; + bool data_valid_{false}; }; //! Constructor @@ -115,10 +122,10 @@ class BinaryConverterHolder : public IndexHolder { //! Create a new iterator IndexHolder::Iterator::Pointer create_iterator(void) override { IndexHolder::Iterator::Pointer iter = front_->create_iterator(); - return iter - ? IndexHolder::Iterator::Pointer( - new BinaryConverterHolder::Iterator(this, std::move(iter))) - : IndexHolder::Iterator::Pointer(); + return iter ? IndexHolder::Iterator::Pointer( + new BinaryConverterHolder::Iterator( + front_->dimension(), quantizer_, std::move(iter))) + : IndexHolder::Iterator::Pointer(); } private: @@ -156,6 +163,11 @@ class BinaryConverter : public IndexConverter { dimension_ = meta_.dimension(); + quantizer_.reset(new (std::nothrow) ailego::BinaryQuantizer()); + if (!quantizer_) { + return IndexError_NoMemory; + } + size_t dim = ailego::BinaryQuantizer::EncodedSizeInBinary32(dimension_) * 32u; @@ -168,6 +180,7 @@ class BinaryConverter : public IndexConverter { //! Cleanup Converter int cleanup(void) override { + holder_.reset(); return 0; } @@ -228,4 +241,4 @@ class BinaryConverter : public IndexConverter { INDEX_FACTORY_REGISTER_CONVERTER(BinaryConverter); } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/quantizer/binary_reformer.cc b/src/core/quantizer/binary_reformer.cc index 78ddffd31..41a051b4e 100644 --- a/src/core/quantizer/binary_reformer.cc +++ b/src/core/quantizer/binary_reformer.cc @@ -11,6 +11,11 @@ // 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 @@ -53,14 +58,20 @@ class BinaryReformer : public IndexReformer { return IndexError_Unsupported; } - size_t dim = - ailego::BinaryQuantizer::EncodedSizeInBinary32(qmeta.dimension()) * 32u; - out->resize( - IndexMeta::ElementSizeof(IndexMeta::DataType::DT_BINARY32, dim)); + const size_t encoded_words = + ailego::BinaryQuantizer::EncodedSizeInBinary32(qmeta.dimension()); + const size_t dim = encoded_words * 32u; const float *vec = reinterpret_cast(query); - - quantizer_.encode(vec, qmeta.dimension(), - reinterpret_cast(&(*out)[0])); + try { + std::vector encoded(encoded_words); + quantizer_.encode(vec, qmeta.dimension(), encoded.data()); + out->assign(reinterpret_cast(encoded.data()), + encoded.size() * sizeof(uint32_t)); + } catch (const std::bad_alloc &) { + return IndexError_NoMemory; + } catch (const std::length_error &) { + return IndexError_InvalidLength; + } *ometa = qmeta; ometa->set_meta(IndexMeta::DataType::DT_BINARY32, dim); @@ -77,14 +88,33 @@ class BinaryReformer : public IndexReformer { return IndexError_Unsupported; } - size_t dim = - ailego::BinaryQuantizer::EncodedSizeInBinary32(qmeta.dimension()) * 32u; - out->resize(count * IndexMeta::ElementSizeof( - IndexMeta::DataType::DT_BINARY32, dim)); + const size_t encoded_words = + ailego::BinaryQuantizer::EncodedSizeInBinary32(qmeta.dimension()); + const size_t dim = encoded_words * 32u; + if (encoded_words > + (std::numeric_limits::max)() / sizeof(uint32_t)) { + return IndexError_InvalidLength; + } + const size_t encoded_size = encoded_words * sizeof(uint32_t); + if (encoded_size != 0 && + count > (std::numeric_limits::max)() / encoded_size) { + return IndexError_InvalidLength; + } const float *vec = reinterpret_cast(query); - - quantizer_.encode(vec, qmeta.dimension() * count, - reinterpret_cast(&(*out)[0])); + try { + out->resize(static_cast(count) * encoded_size); + std::vector encoded(encoded_words); + for (uint32_t i = 0; i < count; ++i) { + quantizer_.encode(vec + static_cast(i) * qmeta.dimension(), + qmeta.dimension(), encoded.data()); + std::memcpy(out->data() + static_cast(i) * encoded_size, + encoded.data(), encoded_size); + } + } catch (const std::bad_alloc &) { + return IndexError_NoMemory; + } catch (const std::length_error &) { + return IndexError_InvalidLength; + } *ometa = qmeta; ometa->set_meta(IndexMeta::DataType::DT_BINARY32, dim); @@ -105,4 +135,4 @@ class BinaryReformer : public IndexReformer { INDEX_FACTORY_REGISTER_REFORMER(BinaryReformer); } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/quantizer/cosine_converter.cc b/src/core/quantizer/cosine_converter.cc index 907682867..929dcc7fa 100644 --- a/src/core/quantizer/cosine_converter.cc +++ b/src/core/quantizer/cosine_converter.cc @@ -69,6 +69,9 @@ class CosineConverterHolder : public IndexHolder { //! Retrieve pointer of data const void *data(void) const override { + if (!data_valid_) { + return nullptr; + } return type_ == original_type_ ? normalize_buffer_.data() : buffer_.data(); } @@ -92,9 +95,14 @@ class CosineConverterHolder : public IndexHolder { private: //! Encode the data by quantizer void convert_record(void) { + data_valid_ = false; if (!front_iter_->is_valid()) { return; } + const void *source_data = front_iter_->data(); + if (!source_data) { + return; + } size_t element_size = owner_->element_size(); size_t original_element_size = @@ -102,8 +110,7 @@ class CosineConverterHolder : public IndexHolder { if (original_type_ == IndexMeta::DataType::DT_FP16) { ::memcpy(reinterpret_cast(&normalize_buffer_[0]), - reinterpret_cast(front_iter_->data()), - original_element_size); + static_cast(source_data), original_element_size); ailego::Float16 *buf = reinterpret_cast(&normalize_buffer_[0]); @@ -117,8 +124,7 @@ class CosineConverterHolder : public IndexHolder { &norm, NORM_SIZE); } else { // original_type_ == IndexMeta::DataType::DT_FP32 ::memcpy(reinterpret_cast(&normalize_buffer_[0]), - reinterpret_cast(front_iter_->data()), - original_element_size); + static_cast(source_data), original_element_size); float *buf = reinterpret_cast(&normalize_buffer_[0]); const float *vec = buf; @@ -156,6 +162,7 @@ class CosineConverterHolder : public IndexHolder { &norm, NORM_SIZE); } } + data_valid_ = true; } //! Members @@ -168,6 +175,7 @@ class CosineConverterHolder : public IndexHolder { size_t original_dimension_{0u}; IndexMeta::DataType original_type_{IndexMeta::DataType::DT_UNDEFINED}; IndexMeta::DataType type_{IndexMeta::DataType::DT_UNDEFINED}; + bool data_valid_{false}; }; //! Constructor @@ -349,6 +357,7 @@ class CosineConverter : public IndexConverter { //! Cleanup Converter int cleanup(void) override { *stats_.mutable_transformed_count() = 0; + holder_.reset(); return 0; } diff --git a/src/core/quantizer/half_float_converter.cc b/src/core/quantizer/half_float_converter.cc index 078a74a30..6330c2807 100644 --- a/src/core/quantizer/half_float_converter.cc +++ b/src/core/quantizer/half_float_converter.cc @@ -40,7 +40,7 @@ class HalfFloatHolder : public IndexHolder { //! Retrieve pointer of data const void *data(void) const override { - return buffer_.data(); + return data_valid_ ? buffer_.data() : nullptr; } //! Test if the iterator is valid @@ -61,15 +61,22 @@ class HalfFloatHolder : public IndexHolder { private: inline void transform_record(void) { - if (front_iter_->is_valid()) { - ailego::FloatHelper::ToFP16( - reinterpret_cast(front_iter_->data()), - buffer_.size(), buffer_.data()); + data_valid_ = false; + if (!front_iter_->is_valid()) { + return; } + const void *data = front_iter_->data(); + if (!data) { + return; + } + ailego::FloatHelper::ToFP16(reinterpret_cast(data), + buffer_.size(), buffer_.data()); + data_valid_ = true; } std::vector buffer_{}; IndexHolder::Iterator::Pointer front_iter_{}; + bool data_valid_{false}; }; //! Constructor @@ -142,6 +149,7 @@ class HalfFloatConverter : public IndexConverter { //! Cleanup Converter int cleanup(void) override { + holder_.reset(); return 0; } @@ -230,7 +238,7 @@ class HalfFloatSparseHolder : public IndexSparseHolder { //! Retrieve sparse data const void *sparse_data() const override { - return sparse_buffer_.data(); + return data_valid_ ? sparse_buffer_.data() : nullptr; } //! Next iterator @@ -241,17 +249,27 @@ class HalfFloatSparseHolder : public IndexSparseHolder { private: inline void transform_record(void) { - if (front_iter_->is_valid()) { - ailego::FloatHelper::ToFP16( - reinterpret_cast(front_iter_->sparse_data()), - front_iter_->sparse_count(), sparse_buffer_.data()); + data_valid_ = false; + if (!front_iter_->is_valid()) { + return; + } + const uint32_t sparse_count = front_iter_->sparse_count(); + const void *data = front_iter_->sparse_data(); + if (sparse_count != 0 && !data) { + return; + } + if (sparse_count != 0) { + ailego::FloatHelper::ToFP16(reinterpret_cast(data), + sparse_count, sparse_buffer_.data()); } + data_valid_ = true; } constexpr static uint32_t MAX_DIM_COUNT = 4096; std::vector sparse_buffer_{}; IndexSparseHolder::Iterator::Pointer front_iter_{}; + bool data_valid_{false}; }; //! Constructor @@ -319,6 +337,7 @@ class HalfFloatSparseConverter : public IndexConverter { //! Cleanup Converter int cleanup(void) override { + holder_.reset(); return 0; } @@ -367,4 +386,4 @@ INDEX_FACTORY_REGISTER_CONVERTER(HalfFloatConverter); INDEX_FACTORY_REGISTER_CONVERTER(HalfFloatSparseConverter); } // namespace core -} // namespace zvec \ No newline at end of file +} // namespace zvec diff --git a/src/core/quantizer/half_float_reformer.cc b/src/core/quantizer/half_float_reformer.cc index 0803a8598..5aed4b09a 100644 --- a/src/core/quantizer/half_float_reformer.cc +++ b/src/core/quantizer/half_float_reformer.cc @@ -177,7 +177,7 @@ class HalfFloatSparseReformer : public IndexReformer { reinterpret_cast(&(*out)[0])); *ometa = qmeta; - ometa->set_data_type(IndexMeta::DataType::DT_FP16); + ometa->set_meta(IndexMeta::DataType::DT_FP16, 0); break; diff --git a/src/core/quantizer/integer_quantizer_converter.cc b/src/core/quantizer/integer_quantizer_converter.cc index cfba26848..f166b0688 100644 --- a/src/core/quantizer/integer_quantizer_converter.cc +++ b/src/core/quantizer/integer_quantizer_converter.cc @@ -49,7 +49,7 @@ class IntegerQuantizerConverterHolder : public IndexHolder { //! Retrieve pointer of data const void *data(void) const override { - return buffer_.data(); + return data_valid_ ? buffer_.data() : nullptr; } //! Test if the iterator is valid @@ -71,12 +71,18 @@ class IntegerQuantizerConverterHolder : public IndexHolder { private: //! Encode the data by quantizer inline void encode_record(void) { - if (front_iter_->is_valid()) { - const float *vec = reinterpret_cast(front_iter_->data()); - quantizer_->encode( - vec, dim_, - reinterpret_cast(buffer_.data())); + data_valid_ = false; + if (!front_iter_->is_valid()) { + return; } + const float *vec = static_cast(front_iter_->data()); + if (!vec) { + return; + } + quantizer_->encode( + vec, dim_, + reinterpret_cast(buffer_.data())); + data_valid_ = true; } //! Members @@ -84,6 +90,7 @@ class IntegerQuantizerConverterHolder : public IndexHolder { IndexHolder::Iterator::Pointer front_iter_{}; std::shared_ptr quantizer_{}; size_t dim_{0u}; + bool data_valid_{false}; }; //! Constructor @@ -203,6 +210,7 @@ class IntegerQuantizerConverter : public IndexConverter { //! Cleanup Converter int cleanup(void) override { + holder_.reset(); return 0; } @@ -227,7 +235,11 @@ class IntegerQuantizerConverter : public IndexConverter { float max = -std::numeric_limits::max(); float min = std::numeric_limits::max(); for (; iter->is_valid(); iter->next()) { - const float *vec = reinterpret_cast(iter->data()); + const float *vec = static_cast(iter->data()); + if (!vec) { + LOG_ERROR("Failed to read holder data while training quantizer"); + return IndexError_ReadData; + } for (size_t i = 0; i < meta_.dimension(); ++i) { max = std::max(max, vec[i]); min = std::min(min, vec[i]); @@ -243,9 +255,13 @@ class IntegerQuantizerConverter : public IndexConverter { return IndexError_Runtime; } for (; iter->is_valid(); iter->next()) { + const float *vec = static_cast(iter->data()); + if (!vec) { + LOG_ERROR("Failed to read holder data while training quantizer"); + return IndexError_ReadData; + } (*stats_.mutable_trained_count())++; - quantizer_->feed(reinterpret_cast(iter->data()), - meta_.dimension()); + quantizer_->feed(vec, meta_.dimension()); } } } else { @@ -259,7 +275,11 @@ class IntegerQuantizerConverter : public IndexConverter { float max = -std::numeric_limits::max(); float min = std::numeric_limits::max(); for (; iter->is_valid(); iter->next()) { - const float *vec = reinterpret_cast(iter->data()); + const float *vec = static_cast(iter->data()); + if (!vec) { + LOG_ERROR("Failed to read holder data while training quantizer"); + return IndexError_ReadData; + } for (size_t i = 0; i < meta_.dimension(); ++i) { max = std::max(max, vec[i]); min = std::min(min, vec[i]); @@ -429,6 +449,7 @@ class IntegerStreamingConverter : public IndexConverter { //! Cleanup Converter int cleanup(void) override { *stats_.mutable_transformed_count() = 0; + holder_.reset(); return 0; } @@ -510,7 +531,7 @@ class IntegerStreamingConverter : public IndexConverter { //! Retrieve pointer of data const void *data(void) const override { - return buffer_.data(); + return data_valid_ ? buffer_.data() : nullptr; } //! Test if the iterator is valid @@ -532,28 +553,31 @@ class IntegerStreamingConverter : public IndexConverter { private: //! Encode the data by quantizer void encode_record(void) { - if (front_iter_->is_valid()) { - const float *vec = - reinterpret_cast(front_iter_->data()); - size_t dim = owner_->dimension_; - if (owner_->rotator_) { - float *rotate_buf = - reinterpret_cast(rotate_buffer_.data()); - owner_->rotator_->rotate(vec, rotate_buf); - vec = rotate_buf; - } - if (owner_->enable_normalize_) { - float norm = 0.0; - memcpy((void *)normalize_buffer_.data(), vec, dim * sizeof(float)); - ailego::Normalizer::L2((float *)normalize_buffer_.data(), - dim, &norm); - vec = (float *)normalize_buffer_.data(); - } - - RecordQuantizer::quantize_record(vec, dim, owner_->data_type(), - owner_->is_euclidean_, - buffer_.data()); + data_valid_ = false; + if (!front_iter_->is_valid()) { + return; + } + const float *vec = static_cast(front_iter_->data()); + if (!vec) { + return; + } + size_t dim = owner_->dimension_; + if (owner_->rotator_) { + float *rotate_buf = reinterpret_cast(rotate_buffer_.data()); + owner_->rotator_->rotate(vec, rotate_buf); + vec = rotate_buf; } + if (owner_->enable_normalize_) { + float norm = 0.0; + memcpy((void *)normalize_buffer_.data(), vec, dim * sizeof(float)); + ailego::Normalizer::L2((float *)normalize_buffer_.data(), dim, + &norm); + vec = (float *)normalize_buffer_.data(); + } + + RecordQuantizer::quantize_record(vec, dim, owner_->data_type(), + owner_->is_euclidean_, buffer_.data()); + data_valid_ = true; } //! Members @@ -562,6 +586,7 @@ class IntegerStreamingConverter : public IndexConverter { std::string normalize_buffer_{}; std::string rotate_buffer_{}; IndexHolder::Iterator::Pointer front_iter_{}; + bool data_valid_{false}; }; //! Constructor diff --git a/src/core/quantizer/mips_converter.cc b/src/core/quantizer/mips_converter.cc index 26b85e180..c25b324ec 100644 --- a/src/core/quantizer/mips_converter.cc +++ b/src/core/quantizer/mips_converter.cc @@ -90,7 +90,7 @@ class MipsConverterHolder : public IndexHolder { //! Retrieve pointer of data const void *data(void) const override { - return buffer_.data(); + return data_valid_ ? buffer_.data() : nullptr; } //! Test if the iterator is valid @@ -112,11 +112,15 @@ class MipsConverterHolder : public IndexHolder { private: //! Transform the data void transform_data(void) { + data_valid_ = false; if (!front_iter_->is_valid()) { return; } - const float *src = reinterpret_cast(front_iter_->data()); + const float *src = static_cast(front_iter_->data()); + if (!src) { + return; + } float *dst = buffer_.data(); if (!spherical_injection_) { ConvertRepeatedQuadraticInjection(src, buffer_.size() - m_value_, @@ -125,6 +129,7 @@ class MipsConverterHolder : public IndexHolder { ConvertSphericalInjection(src, buffer_.size() - m_value_, u_value_, l2_norm_, dst); } + data_valid_ = true; } std::vector buffer_{}; @@ -133,6 +138,7 @@ class MipsConverterHolder : public IndexHolder { float l2_norm_{0.0f}; bool spherical_injection_{false}; IndexHolder::Iterator::Pointer front_iter_{}; + bool data_valid_{false}; }; //! Constructor @@ -218,7 +224,7 @@ class MipsConverterForcedHalfHolder : public IndexHolder { //! Retrieve pointer of data const void *data(void) const override { - return buffer_.data(); + return data_valid_ ? buffer_.data() : nullptr; } //! Test if the iterator is valid @@ -239,11 +245,15 @@ class MipsConverterForcedHalfHolder : public IndexHolder { private: void transform_record(void) { + data_valid_ = false; if (!front_iter_->is_valid()) { return; } - const float *src = reinterpret_cast(front_iter_->data()); + const float *src = static_cast(front_iter_->data()); + if (!src) { + return; + } ailego::Float16 *dst = buffer_.data(); if (!spherical_injection_) { ConvertRepeatedQuadraticInjection(src, buffer_.size() - m_value_, @@ -252,6 +262,7 @@ class MipsConverterForcedHalfHolder : public IndexHolder { ConvertSphericalInjection(src, buffer_.size() - m_value_, u_value_, l2_norm_, dst); } + data_valid_ = true; } std::vector buffer_{}; @@ -260,6 +271,7 @@ class MipsConverterForcedHalfHolder : public IndexHolder { float l2_norm_{0.0f}; bool spherical_injection_{false}; IndexHolder::Iterator::Pointer front_iter_{}; + bool data_valid_{false}; }; //! Constructor @@ -347,7 +359,7 @@ class MipsConverterHalfHolder : public IndexHolder { //! Retrieve pointer of data const void *data(void) const override { - return buffer_.data(); + return data_valid_ ? buffer_.data() : nullptr; } //! Test if the iterator is valid @@ -368,12 +380,16 @@ class MipsConverterHalfHolder : public IndexHolder { private: void transform_record(void) { + data_valid_ = false; if (!front_iter_->is_valid()) { return; } const ailego::Float16 *src = - reinterpret_cast(front_iter_->data()); + static_cast(front_iter_->data()); + if (!src) { + return; + } ailego::Float16 *dst = buffer_.data(); if (!spherical_injection_) { ConvertRepeatedQuadraticInjection(src, buffer_.size() - m_value_, @@ -382,6 +398,7 @@ class MipsConverterHalfHolder : public IndexHolder { ConvertSphericalInjection(src, buffer_.size() - m_value_, u_value_, l2_norm_, dst); } + data_valid_ = true; } std::vector buffer_{}; @@ -390,6 +407,7 @@ class MipsConverterHalfHolder : public IndexHolder { float l2_norm_{0.0f}; bool spherical_injection_{false}; IndexHolder::Iterator::Pointer front_iter_{}; + bool data_valid_{false}; }; //! Constructor @@ -512,6 +530,7 @@ class MipsConverter : public IndexConverter { //! Cleanup Converter int cleanup(void) override { + holder_.reset(); return 0; } @@ -532,10 +551,14 @@ class MipsConverter : public IndexConverter { switch (holder->data_type()) { case IndexMeta::DataType::DT_FP16: for (; iter->is_valid(); iter->next()) { + const auto *vector = + static_cast(iter->data()); + if (!vector) { + LOG_ERROR("Failed to read holder data while training MIPS"); + return IndexError_ReadData; + } float score = 0.0f; - ailego::Norm2Matrix::Compute( - reinterpret_cast(iter->data()), dim, - &score); + ailego::Norm2Matrix::Compute(vector, dim, &score); if (score > l2_norm_) { l2_norm_ = score; @@ -549,9 +572,13 @@ class MipsConverter : public IndexConverter { case IndexMeta::DataType::DT_FP32: for (; iter->is_valid(); iter->next()) { + const auto *vector = static_cast(iter->data()); + if (!vector) { + LOG_ERROR("Failed to read holder data while training MIPS"); + return IndexError_ReadData; + } float score = 0.0f; - ailego::Norm2Matrix::Compute( - reinterpret_cast(iter->data()), dim, &score); + ailego::Norm2Matrix::Compute(vector, dim, &score); if (score > l2_norm_) { l2_norm_ = score; diff --git a/src/core/quantizer/uniform_uint7_converter.cc b/src/core/quantizer/uniform_uint7_converter.cc index b58186707..e7ed59017 100644 --- a/src/core/quantizer/uniform_uint7_converter.cc +++ b/src/core/quantizer/uniform_uint7_converter.cc @@ -130,7 +130,11 @@ class UniformUint7Converter : public IndexConverter { bool all_integer = true; for (; iter->is_valid(); iter->next()) { - const float *vec = reinterpret_cast(iter->data()); + const float *vec = static_cast(iter->data()); + if (!vec) { + LOG_ERROR("UniformUint7Converter: failed to read training vector"); + return IndexError_ReadData; + } for (size_t i = 0; i < original_dimension_; ++i) { float v = vec[i]; if (!std::isfinite(v)) { @@ -254,7 +258,7 @@ class UniformUint7Converter : public IndexConverter { ~Iterator(void) override {} const void *data(void) const override { - return buffer_.data(); + return data_valid_ ? buffer_.data() : nullptr; } bool is_valid(void) const override { @@ -272,10 +276,14 @@ class UniformUint7Converter : public IndexConverter { private: void encode_record(void) { + data_valid_ = false; if (!front_iter_->is_valid()) { return; } - const float *vec = reinterpret_cast(front_iter_->data()); + const float *vec = static_cast(front_iter_->data()); + if (!vec) { + return; + } int8_t *out = buffer_.data(); const float scale = owner_->scale_; const float bias = owner_->bias_; @@ -283,6 +291,7 @@ class UniformUint7Converter : public IndexConverter { if (owner_->quantize_func_ != nullptr) { owner_->quantize_func_(vec, dim, scale, bias, out); + data_valid_ = true; return; } for (size_t i = 0; i < dim; ++i) { @@ -292,11 +301,13 @@ class UniformUint7Converter : public IndexConverter { v = std::max(0.0f, std::min(127.0f, v)); out[i] = static_cast(v); } + data_valid_ = true; } const UniformUint7Holder *owner_{nullptr}; std::vector buffer_{}; IndexHolder::Iterator::Pointer front_iter_{}; + bool data_valid_{false}; }; UniformUint7Holder(IndexHolder::Pointer front, size_t original_dim, diff --git a/src/core/quantizer/uniform_uint8_converter.cc b/src/core/quantizer/uniform_uint8_converter.cc index f78905bc0..6ec6724dc 100644 --- a/src/core/quantizer/uniform_uint8_converter.cc +++ b/src/core/quantizer/uniform_uint8_converter.cc @@ -109,6 +109,10 @@ class UniformUint8Converter : public IndexConverter { for (; iterator->is_valid(); iterator->next()) { const auto *vector = static_cast(iterator->data()); + if (!vector) { + LOG_ERROR("UniformUint8Converter: failed to read training vector"); + return IndexError_ReadData; + } for (size_t i = 0; i < original_dimension_; ++i) { const float value = vector[i]; if (!std::isfinite(value)) { @@ -238,7 +242,7 @@ class UniformUint8Converter : public IndexConverter { } const void *data(void) const override { - return buffer_.data(); + return data_valid_ ? buffer_.data() : nullptr; } bool is_valid(void) const override { @@ -256,17 +260,23 @@ class UniformUint8Converter : public IndexConverter { private: void encode() { + data_valid_ = false; if (!is_valid()) { return; } - EncodeRecord(static_cast(iterator_->data()), - owner_->original_dimension_, owner_->scale_, owner_->bias_, - buffer_.data()); + const float *vector = static_cast(iterator_->data()); + if (!vector) { + return; + } + EncodeRecord(vector, owner_->original_dimension_, owner_->scale_, + owner_->bias_, buffer_.data()); + data_valid_ = true; } const UniformUint8Holder *owner_; std::vector buffer_; IndexHolder::Iterator::Pointer iterator_; + bool data_valid_{false}; }; UniformUint8Holder(IndexHolder::Pointer holder, size_t original_dimension, diff --git a/src/core/utility/file_dumper.cc b/src/core/utility/file_dumper.cc index 070a08eb4..70e03d604 100644 --- a/src/core/utility/file_dumper.cc +++ b/src/core/utility/file_dumper.cc @@ -108,17 +108,22 @@ struct FileDumper : public IndexDumper { //! Close index file bool close_index(void) { if (file_.is_valid()) { + bool succeeded = true; auto write_data = [this](const void *buf, size_t size) { return this->file_.write(buf, size); }; if (!packer_.finish(write_data, stab_)) { LOG_ERROR("Failed to finish packing index package"); - return false; + succeeded = false; + } else if (!file_.flush()) { + LOG_ERROR("Failed to flush packed index file"); + succeeded = false; } stab_.clear(); file_.close(); packer_.reset(); + return succeeded; } return true; } diff --git a/src/core/utility/file_read_storage.cc b/src/core/utility/file_read_storage.cc index a3abb4e5c..cf9c1e478 100644 --- a/src/core/utility/file_read_storage.cc +++ b/src/core/utility/file_read_storage.cc @@ -11,6 +11,7 @@ // 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 @@ -89,10 +90,10 @@ class FileReadStorage : public IndexStorage { //! Fetch data from segment (with own buffer) size_t fetch(size_t offset, void *buf, size_t len) const override { - if (ailego_unlikely(offset + len > region_size_)) { - if (offset > region_size_) { - offset = region_size_; - } + if (ailego_unlikely(offset > region_size_)) { + offset = region_size_; + len = 0; + } else if (ailego_unlikely(len > region_size_ - offset)) { len = region_size_ - offset; } return file_ptr_->read(data_offset_ + offset, buf, len); @@ -100,25 +101,25 @@ class FileReadStorage : public IndexStorage { //! Read data from segment size_t read(size_t offset, const void **data, size_t len) override { - if (ailego_unlikely(offset + len > region_size_)) { - if (offset > region_size_) { - offset = region_size_; - } + if (ailego_unlikely(offset > region_size_)) { + offset = region_size_; + len = 0; + } else if (ailego_unlikely(len > region_size_ - offset)) { len = region_size_ - offset; } - buffer_.reserve(len); + buffer_.resize(len); *data = buffer_.data(); return file_ptr_->read(data_offset_ + offset, (void *)*data, len); } size_t read(size_t offset, MemoryBlock &data, size_t len) override { - if (ailego_unlikely(offset + len > region_size_)) { - if (offset > region_size_) { - offset = region_size_; - } + if (ailego_unlikely(offset > region_size_)) { + offset = region_size_; + len = 0; + } else if (ailego_unlikely(len > region_size_ - offset)) { len = region_size_ - offset; } - buffer_.reserve(len); + buffer_.resize(len); data.reset(buffer_.data()); return file_ptr_->read(data_offset_ + offset, (void *)data.data(), len); } @@ -127,12 +128,15 @@ class FileReadStorage : public IndexStorage { bool read(SegmentData *iovec, size_t count) override { size_t total = 0u; for (auto *it = iovec, *end = iovec + count; it != end; ++it) { - ailego_false_if_false(it->offset + it->length <= region_size_); + ailego_false_if_false(it->offset <= region_size_); + ailego_false_if_false(it->length <= region_size_ - it->offset); + ailego_false_if_false(it->length <= + std::numeric_limits::max() - total); total += it->length; } ailego_false_if_false(total != 0); - buffer_.reserve(total); + buffer_.resize(total); uint8_t *buf = buffer_.data(); for (auto *it = iovec, *end = iovec + count; it != end; ++it) { ailego_false_if_false(file_ptr_->read(data_offset_ + it->offset, buf, @@ -209,10 +213,10 @@ class FileReadStorage : public IndexStorage { //! Fetch data from segment (with own buffer) size_t fetch(size_t offset, void *buf, size_t len) const override { - if (ailego_unlikely(offset + len > region_size_)) { - if (offset > region_size_) { - offset = region_size_; - } + if (ailego_unlikely(offset > region_size_)) { + offset = region_size_; + len = 0; + } else if (ailego_unlikely(len > region_size_ - offset)) { len = region_size_ - offset; } memcpy(buf, data_ + offset, len); @@ -221,10 +225,10 @@ class FileReadStorage : public IndexStorage { //! Read data from segment size_t read(size_t offset, const void **data, size_t len) override { - if (ailego_unlikely(offset + len > region_size_)) { - if (offset > region_size_) { - offset = region_size_; - } + if (ailego_unlikely(offset > region_size_)) { + offset = region_size_; + len = 0; + } else if (ailego_unlikely(len > region_size_ - offset)) { len = region_size_ - offset; } *data = data_ + offset; @@ -232,10 +236,10 @@ class FileReadStorage : public IndexStorage { } size_t read(size_t offset, MemoryBlock &data, size_t len) override { - if (ailego_unlikely(offset + len > region_size_)) { - if (offset > region_size_) { - offset = region_size_; - } + if (ailego_unlikely(offset > region_size_)) { + offset = region_size_; + len = 0; + } else if (ailego_unlikely(len > region_size_ - offset)) { len = region_size_ - offset; } data.reset((void *)(data_ + offset)); @@ -245,7 +249,8 @@ class FileReadStorage : public IndexStorage { //! Read data from segment bool read(SegmentData *iovec, size_t count) override { for (auto *it = iovec, *end = iovec + count; it != end; ++it) { - ailego_false_if_false(it->offset + it->length <= region_size_); + ailego_false_if_false(it->offset <= region_size_); + ailego_false_if_false(it->length <= region_size_ - it->offset); it->data = data_ + it->offset; } return true; @@ -327,15 +332,17 @@ class FileReadStorage : public IndexStorage { size_t size = end_offset > index_offset_ ? end_offset - index_offset_ : 0; auto read_data = [this, &file_ptr, end_offset]( size_t offset, const void **data, size_t len) { - buffer_.reserve(len); - *data = buffer_.data(); - size_t off = index_offset_ + offset; - if (off + len > end_offset) { - if (off > end_offset) { - off = end_offset; - } + size_t off = end_offset; + if (index_offset_ <= end_offset && offset <= end_offset - index_offset_) { + off = index_offset_ + offset; + } else { + len = 0; + } + if (len > end_offset - off) { len = end_offset - off; } + buffer_.resize(len); + *data = buffer_.data(); return file_ptr->read(off, (void *)*data, len); }; diff --git a/src/db/index/column/vector_column/engine_helper.hpp b/src/db/index/column/vector_column/engine_helper.hpp index 59480b7ac..d2cd9850d 100644 --- a/src/db/index/column/vector_column/engine_helper.hpp +++ b/src/db/index/column/vector_column/engine_helper.hpp @@ -250,6 +250,14 @@ class ProximaEngineHelper { auto db_diskann_query_params = dynamic_cast( query_params.query_params.get()); + if (db_diskann_query_params == nullptr) { + return tl::make_unexpected(Status::InvalidArgument( + "DISKANN index requires DiskAnnQueryParams")); + } + if (db_diskann_query_params->list_size() <= 0) { + return tl::make_unexpected(Status::InvalidArgument( + "DiskAnn list_size must be greater than 0")); + } diskann_query_param->list_size = static_cast(db_diskann_query_params->list_size()); } diff --git a/src/db/index/common/query.cc b/src/db/index/common/query.cc index 3f20aa943..02afa293b 100644 --- a/src/db/index/common/query.cc +++ b/src/db/index/common/query.cc @@ -190,6 +190,18 @@ Status QueryTarget::validate(const FieldSchema *schema, "Invalid query: IVF_RABITQ nprobe must be greater than 0"); } } + if (query_params && query_params->type() == IndexType::DISKANN) { + auto diskann_params = + std::dynamic_pointer_cast(query_params); + if (!diskann_params) { + return Status::InvalidArgument( + "Invalid query: DISKANN index requires DiskAnnQueryParams"); + } + if (diskann_params->list_size() <= 0) { + return Status::InvalidArgument( + "Invalid query: DiskAnn list_size must be greater than 0"); + } + } return Status::OK(); } diff --git a/src/db/index/common/schema.cc b/src/db/index/common/schema.cc index d6162dcaa..bbac2323a 100644 --- a/src/db/index/common/schema.cc +++ b/src/db/index/common/schema.cc @@ -218,6 +218,8 @@ Status FieldSchema::validate() const { } if (index_params_->type() == IndexType::DISKANN) { + // DiskAnn supports 64-bit Linux (x86_64/ARM64), macOS ARM64, + // 64-bit Android/iOS, and Windows x86_64. // 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 @@ -227,15 +229,33 @@ Status FieldSchema::validate() const { // // On Linux, DiskAnn prefers io_uring, then libaio, and falls back to // synchronous pread() if neither async backend is available. On macOS, - // DiskAnn uses synchronous pread(). + // Android and iOS, DiskAnn uses synchronous pread(); Windows uses + // overlapped I/O. #if !DISKANN_SUPPORTED return Status::NotSupported( "DiskAnn is not supported on this platform. It is available on " - "Linux (x86_64/ARM64) and macOS (ARM64)."); + "64-bit Linux (x86_64/ARM64), macOS (ARM64), 64-bit Android/iOS, " + "and Windows (x86_64)."); #endif + if (data_type_ != DataType::VECTOR_FP32 && + data_type_ != DataType::VECTOR_FP16) { + return Status::InvalidArgument( + "schema validate failed: DiskAnn only supports FP32/FP16 " + "vector data types"); + } + const auto diskann_quantize_type = vector_index_params->quantize_type(); + if (diskann_quantize_type != QuantizeType::UNDEFINED && + diskann_quantize_type != QuantizeType::FP16) { + return Status::InvalidArgument( + "schema validate failed: DiskAnn only supports FP16 " + "quantization"); + } + if (vector_index_params->quantizer_param().enable_rotate()) { + return Status::InvalidArgument( + "schema validate failed: DiskAnn does not support quantizer " + "rotation"); + } } - - if (vector_index_params->quantize_type() != QuantizeType::UNDEFINED) { auto iter = quantize_type_map.find(data_type_); if (iter == quantize_type_map.end()) { diff --git a/src/include/zvec/ailego/io/io_backend.h b/src/include/zvec/ailego/io/io_backend.h index 9be3dbc69..667193c46 100644 --- a/src/include/zvec/ailego/io/io_backend.h +++ b/src/include/zvec/ailego/io/io_backend.h @@ -24,6 +24,7 @@ #include #include +#include namespace zvec { namespace ailego { @@ -31,23 +32,30 @@ namespace ailego { // Supported DiskAnn I/O backend types. // // Numeric values are part of the C ABI (see zvec_io_backend_type_t in c_api.h): -// kPread = 0, kLibAio = 1, kIoUring = 2. +// kPread = 0, kLibAio = 1, kIoUring = 2, kWindowsOverlapped = 3, +// kUnavailable = 4. enum class IOBackendType { - kPread = 0, // Synchronous pread(); no async I/O - kLibAio = 1, // libaio loaded at runtime via dlopen() - kIoUring = 2, // io_uring via raw kernel syscalls (zero dependency) + kPread = 0, // Synchronous pread(); no async I/O + kLibAio = 1, // libaio loaded at runtime via dlopen() + kIoUring = 2, // io_uring via raw kernel syscalls + kWindowsOverlapped = 3, // Windows overlapped I/O using per-context IOCP + kUnavailable = 4, // DiskAnn is disabled on this target }; -// Returns the currently active I/O backend type. +// Returns the currently selected I/O backend type for new contexts. // Triggers backend selection on first call. Linux tries io_uring, then libaio, -// and finally synchronous pread. macOS ARM64 uses synchronous pread. -IOBackendType current_io_backend_type(); - -// Returns a human-readable description of the currently active I/O backend. +// and finally synchronous pread. If a context cannot initialize the preferred +// Linux backend, the process-wide selection is downgraded so later calls and +// contexts report and use the effective fallback. macOS ARM64 uses synchronous +// pread. Windows uses unbuffered overlapped I/O with a per-context completion +// port. +ZVEC_AILEGO_API IOBackendType current_io_backend_type(); + +// Returns a human-readable description of the currently selected I/O backend. // The description identifies io_uring, libaio, or pread. On Linux, the pread // description also explains that io_uring and libaio were unavailable and // provides guidance for enabling an asynchronous backend. -std::string current_io_backend_description(); +ZVEC_AILEGO_API std::string current_io_backend_description(); } // namespace ailego } // namespace zvec diff --git a/src/include/zvec/c_api.h b/src/include/zvec/c_api.h index 5a6cbb906..67e52d88b 100644 --- a/src/include/zvec/c_api.h +++ b/src/include/zvec/c_api.h @@ -789,17 +789,25 @@ typedef uint32_t zvec_io_backend_type_t; #define ZVEC_IO_BACKEND_TYPE_LIBAIO \ 1 /**< libaio loaded at runtime via dlopen() */ #define ZVEC_IO_BACKEND_TYPE_IO_URING \ - 2 /**< io_uring via raw kernel syscalls (zero dependency) */ + 2 /**< io_uring via raw Linux kernel syscalls (zero dependency) */ +#define ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED \ + 3 /**< Windows overlapped I/O using per-context IOCP */ +#define ZVEC_IO_BACKEND_TYPE_UNAVAILABLE \ + 4 /**< DiskAnn is disabled on this target */ /** * @brief Get the current I/O backend type for DiskAnn disk reads. * * Linux selects the first usable backend in this order: io_uring, libaio, - * then synchronous pread. macOS ARM64 uses synchronous pread. + * then synchronous pread. macOS ARM64 uses synchronous pread. Windows uses + * unbuffered overlapped I/O with a per-context completion port. * * @return zvec_io_backend_type_t The loaded backend type - * (ZVEC_IO_BACKEND_TYPE_IO_URING, ZVEC_IO_BACKEND_TYPE_LIBAIO, - * or ZVEC_IO_BACKEND_TYPE_PREAD). + * ZVEC_IO_BACKEND_TYPE_IO_URING, ZVEC_IO_BACKEND_TYPE_LIBAIO, + * ZVEC_IO_BACKEND_TYPE_PREAD, or + * ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED. Returns + * ZVEC_IO_BACKEND_TYPE_UNAVAILABLE when DiskAnn is disabled for the + * current target architecture. */ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void); @@ -808,7 +816,8 @@ ZVEC_EXPORT zvec_io_backend_type_t ZVEC_CALL zvec_get_io_backend_type(void); * * @param type The backend type code. * @return Thread-local string valid until the next call on this thread; - * "io_uring", "libaio", "pread", or "unknown". + * "io_uring", "libaio", "pread", "windows_overlapped", or + * "unavailable"; unknown numeric values return "unknown". */ ZVEC_EXPORT const char *ZVEC_CALL zvec_get_io_backend_type_name(zvec_io_backend_type_t type); @@ -818,7 +827,8 @@ zvec_get_io_backend_type_name(zvec_io_backend_type_t type); * * The description identifies io_uring, libaio, or pread. On Linux, the pread * description also explains that io_uring and libaio were unavailable and - * provides guidance for enabling an asynchronous backend. + * provides guidance for enabling an asynchronous backend. Windows reports its + * overlapped-I/O backend. * * @return Thread-local string valid until the next call on this thread. */ diff --git a/src/include/zvec/core/framework/index_context.h b/src/include/zvec/core/framework/index_context.h index 618005a1e..b28d3e87b 100644 --- a/src/include/zvec/core/framework/index_context.h +++ b/src/include/zvec/core/framework/index_context.h @@ -301,6 +301,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/src/include/zvec/core/interface/index.h b/src/include/zvec/core/interface/index.h index 3086ac481..b2c79ec8d 100644 --- a/src/include/zvec/core/interface/index.h +++ b/src/include/zvec/core/interface/index.h @@ -14,10 +14,13 @@ #pragma once +#include #include +#include #include #include #include +#include #include #include #include @@ -149,7 +152,7 @@ class ZVEC_CORE_API Index { bool is_dirty() const; - uint32_t get_doc_count() const; + virtual uint32_t get_doc_count() const; core::IndexStreamer::Pointer index_searcher(); @@ -209,8 +212,28 @@ class ZVEC_CORE_API Index { bool init_context(); core::IndexContext::Pointer &acquire_context(); + core::IndexStreamer::Pointer streamer_snapshot() const { + return std::atomic_load_explicit(&streamer_, std::memory_order_acquire); + } + + core::IndexStreamer::Pointer exchange_streamer( + core::IndexStreamer::Pointer replacement) { + return std::atomic_exchange_explicit(&streamer_, std::move(replacement), + std::memory_order_acq_rel); + } + + core::IndexStorage::Pointer storage_snapshot() const { + return std::atomic_load_explicit(&storage_, std::memory_order_acquire); + } + + core::IndexStorage::Pointer exchange_storage( + core::IndexStorage::Pointer replacement) { + return std::atomic_exchange_explicit(&storage_, std::move(replacement), + std::memory_order_acq_rel); + } + protected: - bool is_trained_{false}; + std::atomic is_trained_{false}; BaseIndexParam param_; ailego::Params proxima_index_params_{}; @@ -387,6 +410,8 @@ class ZVEC_CORE_API DiskAnnIndex : public Index { public: DiskAnnIndex() = default; + uint32_t get_doc_count() const override; + protected: int CreateAndInitStreamer(const BaseIndexParam ¶m) override; @@ -408,9 +433,12 @@ class ZVEC_CORE_API DiskAnnIndex : public Index { int GenerateHolder(); private: + int CommitBuiltSnapshot(bool *snapshot_replaced); + DiskAnnIndexParam param_{}; - std::mutex mutex_{}; - std::vector> doc_cache_; + mutable std::mutex mutex_{}; + bool is_training_{false}; + std::map doc_cache_; core::IndexHolder::Pointer holder_{}; std::string file_path_; }; diff --git a/tests/ailego/io/io_backend_test.cc b/tests/ailego/io/io_backend_test.cc index 897f0a9de..cc2133991 100644 --- a/tests/ailego/io/io_backend_test.cc +++ b/tests/ailego/io/io_backend_test.cc @@ -37,6 +37,11 @@ TEST(IOBackend, ConcurrentProbeReturnsStableType) { for (IOBackendType type : results) { EXPECT_EQ(type, results[0]); } +#if !defined(DISKANN_SUPPORTED) || !DISKANN_SUPPORTED + EXPECT_EQ(results[0], IOBackendType::kUnavailable); +#elif defined(_WIN32) || defined(_WIN64) + EXPECT_EQ(results[0], IOBackendType::kWindowsOverlapped); +#endif std::string description = current_io_backend_description(); EXPECT_FALSE(description.empty()); const char *backend_name = ""; @@ -50,6 +55,12 @@ TEST(IOBackend, ConcurrentProbeReturnsStableType) { case IOBackendType::kPread: backend_name = "pread"; break; + case IOBackendType::kWindowsOverlapped: + backend_name = "windows_overlapped"; + break; + case IOBackendType::kUnavailable: + backend_name = "unavailable"; + break; } EXPECT_NE(description.find(backend_name), std::string::npos); } diff --git a/tests/ailego/parallel/thread_queue_test.cc b/tests/ailego/parallel/thread_queue_test.cc index 1a1fff39a..a1fbd1d9f 100644 --- a/tests/ailego/parallel/thread_queue_test.cc +++ b/tests/ailego/parallel/thread_queue_test.cc @@ -56,7 +56,7 @@ TEST(ThreadQueue, General) { { std::unique_lock lock(count_mutex); completed = count_cond.wait_for(lock, std::chrono::seconds(10), - [&count]() { return count == kTaskCount; }); + [&]() { return count == kTaskCount; }); completed_count = count; } @@ -98,7 +98,7 @@ TEST(ThreadQueue, MutliThread) { { std::unique_lock lock(count_mutex); completed = count_cond.wait_for(lock, std::chrono::seconds(10), - [&count]() { return count == kTaskCount; }); + [&]() { return count == kTaskCount; }); completed_count = count; } 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 02cb6804e..8061ed254 100644 --- a/tests/c/c_api_test.c +++ b/tests/c/c_api_test.c @@ -162,10 +162,20 @@ void test_io_backend_functions(void) { TEST_ASSERT( strcmp(zvec_get_io_backend_type_name(ZVEC_IO_BACKEND_TYPE_IO_URING), "io_uring") == 0); + TEST_ASSERT(strcmp(zvec_get_io_backend_type_name( + ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED), + "windows_overlapped") == 0); + TEST_ASSERT( + strcmp(zvec_get_io_backend_type_name(ZVEC_IO_BACKEND_TYPE_UNAVAILABLE), + "unavailable") == 0); TEST_ASSERT(strcmp(zvec_get_io_backend_type_name(999), "unknown") == 0); zvec_io_backend_type_t current = zvec_get_io_backend_type(); -#if defined(__APPLE__) && defined(__MACH__) +#if !defined(DISKANN_SUPPORTED) || !DISKANN_SUPPORTED + TEST_ASSERT(current == ZVEC_IO_BACKEND_TYPE_UNAVAILABLE); +#elif defined(_WIN32) + TEST_ASSERT(current == ZVEC_IO_BACKEND_TYPE_WINDOWS_OVERLAPPED); +#elif defined(__APPLE__) && defined(__MACH__) TEST_ASSERT(current == ZVEC_IO_BACKEND_TYPE_PREAD); #else TEST_ASSERT(current == ZVEC_IO_BACKEND_TYPE_PREAD || @@ -502,11 +512,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"); @@ -5230,7 +5240,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, @@ -5271,7 +5281,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..8aef65d25 100644 --- a/tests/core/algorithm/diskann/CMakeLists.txt +++ b/tests/core/algorithm/diskann/CMakeLists.txt @@ -2,6 +2,29 @@ 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$") +else() + if(WIN32) + # Windows has a dedicated overlapped-I/O suite; the mobile compatibility + # suite intentionally exercises POSIX file-descriptor behavior. + list(FILTER ALL_TEST_SRCS EXCLUDE REGEX + "/diskann_file_reader_test\\.cc$|/diskann_mobile_compat_test\\.cc$") + else() + list(FILTER ALL_TEST_SRCS EXCLUDE REGEX + "/diskann_file_reader_windows_test\\.cc$") + endif() + + if(NOT CMAKE_SYSTEM_NAME STREQUAL "Linux") + list(FILTER ALL_TEST_SRCS EXCLUDE REGEX + "/diskann_file_reader_aio_test\\.cc$") + endif() +endif() + foreach(CC_SRCS ${ALL_TEST_SRCS}) get_filename_component(CC_TARGET ${CC_SRCS} NAME_WE) cc_gtest( @@ -11,4 +34,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_builder_test.cc b/tests/core/algorithm/diskann/diskann_builder_test.cc index 0f9a8ea7e..9eafddecf 100644 --- a/tests/core/algorithm/diskann/diskann_builder_test.cc +++ b/tests/core/algorithm/diskann/diskann_builder_test.cc @@ -19,10 +19,12 @@ #include #include #include +#include #include #include #include -#include "diskann_holder.h" +#include "diskann_builder_entity.h" +#include "diskann_context.h" #include "diskann_params.h" using namespace zvec::core; @@ -43,6 +45,40 @@ class DiskAnnBuilderTest : public testing::Test { std::string DiskAnnBuilderTest::_dir("DiskAnnBuilderTest"); shared_ptr DiskAnnBuilderTest::_index_meta_ptr; +class CountOverrideHolder : public IndexHolder { + public: + CountOverrideHolder(IndexHolder::Pointer holder, size_t declared_count) + : holder_(std::move(holder)), declared_count_(declared_count) {} + + size_t count() const override { + return declared_count_; + } + + size_t dimension() const override { + return holder_->dimension(); + } + + IndexMeta::DataType data_type() const override { + return holder_->data_type(); + } + + size_t element_size() const override { + return holder_->element_size(); + } + + bool multipass() const override { + return holder_->multipass(); + } + + Iterator::Pointer create_iterator() override { + return holder_->create_iterator(); + } + + private: + IndexHolder::Pointer holder_; + size_t declared_count_; +}; + void DiskAnnBuilderTest::SetUp(void) { LoggerBroker::SetLevel(Logger::LEVEL_INFO); @@ -85,6 +121,11 @@ TEST_F(DiskAnnBuilderTest, TestGeneral) { ASSERT_EQ(0, builder->build(holder)); + // Dump must use the vectors captured by build, not reread a holder that may + // have changed after the graph and PQ data were finalized. + NumericalVector late_vector(dim, -1.0F); + ASSERT_TRUE(holder->emplace(doc_cnt, late_vector)); + auto dumper = IndexFactory::CreateDumper("FileDumper"); ASSERT_NE(dumper, nullptr); @@ -102,6 +143,213 @@ TEST_F(DiskAnnBuilderTest, TestGeneral) { ASSERT_GT(stats.built_costtime(), 0UL); } +TEST_F(DiskAnnBuilderTest, RejectsDuplicateValidKeysAtDump) { + DiskAnnBuilderEntity entity; + ASSERT_EQ(0, entity.init(*_index_meta_ptr, 16, 32, 0.0, 1)); + + std::vector vector(dim, 1.0F); + EXPECT_EQ(0, entity.add_vector(42, vector.data())); + EXPECT_EQ(0, entity.add_vector(42, vector.data())); + + // Invalid keys represent empty/deleted slots and may legitimately repeat. + EXPECT_EQ(0, entity.add_vector(kInvalidKey, vector.data())); + EXPECT_EQ(0, entity.add_vector(kInvalidKey, vector.data())); + EXPECT_EQ(4U, entity.doc_cnt()); + EXPECT_EQ(IndexError_InvalidArgument, entity.add_vector(7, nullptr)); + + auto dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_NE(nullptr, dumper); + ASSERT_EQ(0, dumper->create(_dir + "/DuplicateKeys")); + EXPECT_EQ(IndexError_Exist, entity.dump_key_mapping_segment(dumper)); + EXPECT_EQ(0, dumper->close()); + + DiskAnnBuilderEntity holes; + ASSERT_EQ(0, holes.init(*_index_meta_ptr, 16, 32, 0.0, 1)); + ASSERT_EQ(0, holes.add_vector(kInvalidKey, vector.data())); + ASSERT_EQ(0, holes.add_vector(kInvalidKey, vector.data())); + auto holes_dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_NE(nullptr, holes_dumper); + ASSERT_EQ(0, holes_dumper->create(_dir + "/InvalidKeySlots")); + EXPECT_EQ(0, holes.dump_key_mapping_segment(holes_dumper)); + EXPECT_EQ(0, holes_dumper->close()); +} + +TEST_F(DiskAnnBuilderTest, NeighborStorageIsNaturallyAligned) { + DiskAnnBuilderEntity entity; + ASSERT_EQ(0, entity.init(*_index_meta_ptr, 1, 1, 0.0, 1)); + + std::vector vector(dim, 1.0F); + ASSERT_EQ(0, entity.add_vector(0, vector.data())); + + auto neighbors = entity.get_neighbors(0); + ASSERT_NE(nullptr, neighbors.second); + EXPECT_EQ(0U, reinterpret_cast(neighbors.second) % + alignof(diskann_id_t)); + + ASSERT_EQ(0, entity.add_neighbor(0, 7)); + neighbors = entity.get_neighbors(0); + ASSERT_EQ(1U, neighbors.first); + EXPECT_EQ(7U, neighbors.second[0]); + + DiskAnnBuilderEntity slack_entity; + ASSERT_EQ(0, slack_entity.init(*_index_meta_ptr, 100, 100, 0.0, 1)); + ASSERT_EQ(0, slack_entity.add_vector(0, vector.data())); + for (diskann_id_t neighbor_id = 0; neighbor_id < 130; ++neighbor_id) { + ASSERT_EQ(0, slack_entity.add_neighbor(0, neighbor_id)); + } + EXPECT_EQ(IndexError_IndexFull, slack_entity.add_neighbor(0, 130)); +} + +TEST_F(DiskAnnBuilderTest, RejectsInvalidGraphAndSamplingParameters) { + auto expect_invalid = [](const Params ¶ms) { + auto builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); + ASSERT_NE(nullptr, builder); + EXPECT_EQ(IndexError_InvalidArgument, + builder->init(*_index_meta_ptr, params)); + }; + + Params params; + params.set(PARAM_DISKANN_BUILDER_MAX_DEGREE, 0U); + expect_invalid(params); + + params.clear(); + params.set(PARAM_DISKANN_BUILDER_LIST_SIZE, 0U); + expect_invalid(params); + + params.clear(); + params.set(PARAM_DISKANN_BUILDER_MAX_TRAIN_SAMPLE_COUNT, 0U); + expect_invalid(params); + + params.clear(); + params.set(PARAM_DISKANN_BUILDER_TRAIN_SAMPLE_RATIO, 0.0); + expect_invalid(params); + + params.clear(); + params.set(PARAM_DISKANN_BUILDER_TRAIN_SAMPLE_RATIO, 1.01); + expect_invalid(params); + + params.clear(); + params.set(PARAM_DISKANN_BUILDER_TRAIN_SAMPLE_RATIO, + std::numeric_limits::quiet_NaN()); + expect_invalid(params); + + DiskAnnBuilderEntity entity; + EXPECT_EQ(IndexError_InvalidArgument, + entity.init(*_index_meta_ptr, 0, 1, 0.0, 1)); + EXPECT_EQ(IndexError_InvalidArgument, + entity.init(*_index_meta_ptr, 1, 0, 0.0, 1)); + ASSERT_EQ(0, entity.init(*_index_meta_ptr, 1, 1, 0.0, 1)); + EXPECT_EQ(IndexError_InvalidLength, entity.reserve_space(0)); + EXPECT_EQ(IndexError_InvalidLength, + entity.reserve_space((std::numeric_limits::max)())); +} + +TEST_F(DiskAnnBuilderTest, RejectsMismatchedHolderMetadata) { + Params params; + params.set(PARAM_DISKANN_BUILDER_MAX_DEGREE, 16); + params.set(PARAM_DISKANN_BUILDER_LIST_SIZE, 20); + params.set(PARAM_DISKANN_BUILDER_MAX_PQ_CHUNK_NUM, 1); + params.set(PARAM_DISKANN_BUILDER_THREAD_COUNT, 1); + + auto builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); + ASSERT_NE(nullptr, builder); + ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); + + auto mismatched = + make_shared>(dim / 2); + NumericalVector short_vector(dim / 2, 1.0F); + ASSERT_TRUE(mismatched->emplace(0, short_vector)); + EXPECT_EQ(IndexError_Mismatch, builder->train(mismatched)); + + auto valid = + make_shared>(dim); + NumericalVector vector(dim, 1.0F); + ASSERT_TRUE(valid->emplace(0, vector)); + ASSERT_EQ(0, builder->train(valid)); + EXPECT_EQ(IndexError_Mismatch, builder->build(mismatched)); + EXPECT_EQ(0, builder->build(valid)); +} + +TEST_F(DiskAnnBuilderTest, FailedBuildDiscardsPartialEntity) { + Params params; + params.set(PARAM_DISKANN_BUILDER_MAX_DEGREE, 16); + params.set(PARAM_DISKANN_BUILDER_LIST_SIZE, 20); + params.set(PARAM_DISKANN_BUILDER_MAX_PQ_CHUNK_NUM, 1); + params.set(PARAM_DISKANN_BUILDER_THREAD_COUNT, 1); + + auto builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); + ASSERT_NE(nullptr, builder); + ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); + + auto valid = + make_shared>(dim); + for (size_t i = 0; i < 2; ++i) { + NumericalVector vector(dim, static_cast(i)); + ASSERT_TRUE(valid->emplace(i, vector)); + } + ASSERT_EQ(0, builder->train(valid)); + + auto incomplete = + make_shared>(dim); + NumericalVector vector(dim, 1.0F); + ASSERT_TRUE(incomplete->emplace(0, vector)); + auto wrong_count = make_shared(incomplete, 2); + EXPECT_EQ(IndexError_InvalidLength, builder->build(wrong_count)); + + // A failed build invalidates the trained state instead of appending a retry + // to the partial entity. Retraining starts from a clean entity. + EXPECT_EQ(IndexError_NoReady, builder->build(valid)); + ASSERT_EQ(0, builder->train(valid)); + EXPECT_EQ(0, builder->build(valid)); +} + +TEST_F(DiskAnnBuilderTest, PqSamplingUsesRatioAndWholeDataset) { + constexpr size_t kSamplingDim = 1; + constexpr size_t kDocCount = 100; + IndexMeta meta(IndexMeta::DataType::DT_FP32, kSamplingDim); + meta.set_metric("SquaredEuclidean", 0, Params()); + auto holder = make_shared>( + kSamplingDim); + for (size_t i = 0; i < kDocCount; ++i) { + NumericalVector vector(kSamplingDim, static_cast(i)); + ASSERT_TRUE(holder->emplace(i, vector)); + } + + DiskAnnPqTrainer trainer(10, 0.5); + std::string sample_data; + size_t sample_size = 0; + ASSERT_EQ(0, + trainer.gen_random_sample(holder, meta, sample_data, sample_size)); + EXPECT_EQ(10U, sample_size); + ASSERT_EQ(sample_size * sizeof(float), sample_data.size()); + + std::string prefix(sample_data.size(), '\0'); + for (size_t i = 0; i < sample_size; ++i) { + const float value = static_cast(i); + std::memcpy(prefix.data() + i * sizeof(float), &value, sizeof(value)); + } + EXPECT_NE(prefix, sample_data) + << "PQ sampling must not always select the first vectors"; + + DiskAnnPqTrainer repeated_trainer(10, 0.5); + std::string repeated_data; + size_t repeated_size = 0; + ASSERT_EQ(0, repeated_trainer.gen_random_sample(holder, meta, repeated_data, + repeated_size)); + EXPECT_EQ(sample_size, repeated_size); + EXPECT_EQ(sample_data, repeated_data); + + DiskAnnPqTrainer ratio_limited_trainer(100, 0.05); + ASSERT_EQ(0, ratio_limited_trainer.gen_random_sample( + holder, meta, sample_data, sample_size)); + EXPECT_EQ(5U, sample_size); + + DiskAnnPqTrainer invalid_trainer(0, 1.0); + EXPECT_EQ(IndexError_InvalidArgument, + invalid_trainer.gen_random_sample(holder, meta, sample_data, + sample_size)); +} + // Regression test: building a small DiskAnn index must complete quickly. // A lost-wakeup bug in the condition-variable progress loops previously caused // 15–30 second stalls during train/build on small datasets because @@ -145,6 +393,55 @@ TEST_F(DiskAnnBuilderTest, SmallDatasetBuildTime) { << " ms — likely a lost-wakeup regression in progress loops."; } +TEST_F(DiskAnnBuilderTest, SingleDocumentIndexCanBeSearched) { + auto builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); + ASSERT_NE(nullptr, builder); + + auto holder = + make_shared>(dim); + NumericalVector vector(dim, 1.0F); + ASSERT_TRUE(holder->emplace(7, vector)); + + Params params; + params.set(PARAM_DISKANN_BUILDER_MAX_DEGREE, 16); + params.set(PARAM_DISKANN_BUILDER_LIST_SIZE, 32); + params.set(PARAM_DISKANN_BUILDER_MAX_PQ_CHUNK_NUM, 1); + params.set(PARAM_DISKANN_BUILDER_THREAD_COUNT, 1); + ASSERT_EQ(0, builder->init(*_index_meta_ptr, params)); + ASSERT_EQ(0, builder->train(holder)); + ASSERT_EQ(0, builder->build(holder)); + + const string path = _dir + "/SingleDocumentIndexCanBeSearched"; + auto dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_NE(nullptr, dumper); + ASSERT_EQ(0, dumper->create(path)); + ASSERT_EQ(0, builder->dump(dumper)); + ASSERT_EQ(0, dumper->close()); + + auto storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->open(path, false)); + + auto searcher = IndexFactory::CreateSearcher("DiskAnnSearcher"); + ASSERT_NE(nullptr, searcher); + Params search_params; + search_params.set(PARAM_DISKANN_SEARCHER_LIST_SIZE, 32); + ASSERT_EQ(0, searcher->init(search_params)); + ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); + + auto context = searcher->create_context(); + ASSERT_NE(nullptr, context); + auto *diskann_context = dynamic_cast(context.get()); + ASSERT_NE(nullptr, diskann_context); + EXPECT_EQ(1U, diskann_context->list_size()); + + context->set_topk(1); + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, dim); + ASSERT_EQ(0, searcher->search_impl(vector.data(), query_meta, context)); + ASSERT_EQ(1U, context->result().size()); + EXPECT_EQ(7U, context->result()[0].key()); +} + TEST_F(DiskAnnBuilderTest, MemoryLimitCapsPqChunkCount) { constexpr size_t kTestDim = 8; constexpr size_t kDocCnt = 16; diff --git a/tests/core/algorithm/diskann/diskann_cache_test.cc b/tests/core/algorithm/diskann/diskann_cache_test.cc new file mode 100644 index 000000000..70c0ad1c3 --- /dev/null +++ b/tests/core/algorithm/diskann/diskann_cache_test.cc @@ -0,0 +1,80 @@ +// 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 "diskann_params.h" +#include "diskann_util.h" + +using namespace zvec::core; + +TEST(DiskAnnCachePreloadTest, BoundsBatchByScratchBufferCapacity) { + EXPECT_EQ(DiskAnnUtil::cache_load_batch_size(0), 1u); + EXPECT_EQ(DiskAnnUtil::cache_load_batch_size(1), 128u); + EXPECT_EQ(DiskAnnUtil::cache_load_batch_size(2), 64u); + EXPECT_EQ(DiskAnnUtil::cache_load_batch_size(128), 1u); + EXPECT_EQ(DiskAnnUtil::cache_load_batch_size(129), 1u); +} + +TEST(DiskAnnCacheConfigTest, FailedReinitKeepsPreviousValidConfiguration) { + zvec::ailego::Params valid; + valid.set(PARAM_DISKANN_SEARCHER_LIST_SIZE, 321); + valid.set(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, 7); + + zvec::ailego::Params invalid; + invalid.set(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, -1); + + auto searcher = IndexFactory::CreateSearcher("DiskAnnSearcher"); + ASSERT_NE(searcher, nullptr); + ASSERT_EQ(0, searcher->init(valid)); + ASSERT_EQ(IndexError_InvalidArgument, searcher->init(invalid)); + uint32_t list_size = 0; + long long cache_nodes = 0; + EXPECT_TRUE( + searcher->params().get(PARAM_DISKANN_SEARCHER_LIST_SIZE, &list_size)); + EXPECT_TRUE(searcher->params().get(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, + &cache_nodes)); + EXPECT_EQ(list_size, 321u); + EXPECT_EQ(cache_nodes, 7); + + auto streamer = IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(streamer, nullptr); + IndexMeta meta(IndexMeta::DataType::DT_FP32, 8); + IndexMeta invalid_meta(IndexMeta::DataType::DT_FP32, 16); + ASSERT_EQ(0, streamer->init(meta, valid)); + ASSERT_EQ(IndexError_InvalidArgument, streamer->init(invalid_meta, invalid)); + EXPECT_EQ(streamer->meta().dimension(), 8u); +} + +TEST(DiskAnnCacheConfigTest, RejectsOutOfRangeNodeCount) { + zvec::ailego::Params negative; + negative.set(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, -1); + + zvec::ailego::Params too_large; + const long long too_many_nodes = + static_cast((std::numeric_limits::max)()) + 1; + too_large.set(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, too_many_nodes); + + auto searcher = IndexFactory::CreateSearcher("DiskAnnSearcher"); + ASSERT_NE(searcher, nullptr); + EXPECT_EQ(IndexError_InvalidArgument, searcher->init(negative)); + EXPECT_EQ(IndexError_InvalidArgument, searcher->init(too_large)); + + auto streamer = IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(streamer, nullptr); + IndexMeta meta(IndexMeta::DataType::DT_FP32, 8); + EXPECT_EQ(IndexError_InvalidArgument, streamer->init(meta, negative)); + EXPECT_EQ(IndexError_InvalidArgument, streamer->init(meta, too_large)); +} diff --git a/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc index a416b0623..bec83670a 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_aio_test.cc @@ -12,10 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include "diskann_file_reader.h" - -#if defined(__linux) || defined(__linux__) - #include #include #include @@ -27,6 +23,7 @@ #include #include #include +#include "diskann_file_reader.h" namespace zvec { namespace core { @@ -257,5 +254,3 @@ TEST(DiskAnnLinuxAioTest, DrainsAllCompletionsBeforePreadFallback) { std::free(output); } - -#endif // __linux__ diff --git a/tests/core/algorithm/diskann/diskann_file_reader_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_test.cc index 81542e079..f1b2d0877 100644 --- a/tests/core/algorithm/diskann/diskann_file_reader_test.cc +++ b/tests/core/algorithm/diskann/diskann_file_reader_test.cc @@ -167,6 +167,83 @@ TEST(DiskAnnFileReaderTest, ReadBeforeOpenReturnsError) { EXPECT_NE(reader.read(requests, ctx, false), 0); } +TEST(DiskAnnFileReaderTest, + OpenFromHandleSurvivesPathReplacementBeforeHandoff) { + TemporaryFile original; + TemporaryFile replacement; + ASSERT_GE(original.fd(), 0); + ASSERT_GE(replacement.fd(), 0); + + std::vector original_data(kPageSize, 0x3a); + std::vector replacement_data(kPageSize, 0xc7); + ASSERT_TRUE(original.write_all(original_data.data(), original_data.size())); + ASSERT_TRUE( + replacement.write_all(replacement_data.data(), replacement_data.size())); + + const int source_flags_before = ::fcntl(original.fd(), F_GETFL); + ASSERT_GE(source_flags_before, 0); + + // The original descriptor represents FileReadStorage after it supplied + // metadata. Replace the path first to prove the handoff follows the open + // file object rather than resolving the path again. + ASSERT_EQ(::rename(replacement.path(), original.path()), 0); + + LinuxAlignedFileReader original_reader; + ASSERT_EQ(original_reader.open_from_handle(original.path(), original.fd()), + 0); + const int source_flags_after = ::fcntl(original.fd(), F_GETFL); + ASSERT_GE(source_flags_after, 0); + EXPECT_EQ(source_flags_after, source_flags_before); + + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector requests{{0, kPageSize, output.get()}}; + IOContext ctx{}; + ASSERT_EQ(setup_io_ctx(ctx), 0); + + ASSERT_EQ(original_reader.read(requests, ctx, false), 0); + EXPECT_EQ(std::memcmp(output.get(), original_data.data(), kPageSize), 0); + + LinuxAlignedFileReader replacement_reader; + replacement_reader.open(original.path()); + ASSERT_EQ(replacement_reader.read(requests, ctx, false), 0); + EXPECT_EQ(std::memcmp(output.get(), replacement_data.data(), kPageSize), 0); + + EXPECT_EQ(destroy_io_ctx(ctx), 0); + original_reader.close(); + replacement_reader.close(); +} + +#if defined(__linux__) || defined(__linux) || defined(__APPLE__) || \ + defined(__MACH__) +TEST(DiskAnnFileReaderTest, OpenFromHandleDoesNotChangeSourceFlags) { + TemporaryFile file; + ASSERT_GE(file.fd(), 0); + + std::vector source(kPageSize, 0x6b); + ASSERT_TRUE(file.write_all(source.data(), source.size())); + + const int flags_before = ::fcntl(file.fd(), F_GETFL); + ASSERT_GE(flags_before, 0); +#if defined(__linux__) || defined(__linux) + ASSERT_EQ(flags_before & O_DIRECT, 0); +#endif + + LinuxAlignedFileReader reader; + ASSERT_EQ(reader.open_from_handle(file.path(), file.fd()), 0); + + const int flags_after = ::fcntl(file.fd(), F_GETFL); + ASSERT_GE(flags_after, 0); + EXPECT_EQ(flags_after, flags_before); + + // An unaligned one-byte read remains valid on the caller's buffered handle. + uint8_t byte = 0; + ASSERT_EQ(::pread(file.fd(), &byte, 1, 0), 1); + EXPECT_EQ(byte, 0x6b); + reader.close(); +} +#endif + #if defined(__APPLE__) || defined(__MACH__) TEST(DiskAnnFileReaderTest, MacOSBatchUsesSynchronousPread) { TemporaryFile file; diff --git a/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc b/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc new file mode 100644 index 000000000..5a515f92a --- /dev/null +++ b/tests/core/algorithm/diskann/diskann_file_reader_windows_test.cc @@ -0,0 +1,730 @@ +// 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 +#include +#include "diskann_file_reader.h" + +namespace zvec { +namespace core { + +class WindowsAlignedFileReaderTestPeer { + public: + static HANDLE stable_file_handle(const WindowsAlignedFileReader &reader) { + return reader.stable_file_handle_; + } +}; + +} // namespace core +} // namespace zvec + +using namespace zvec::core; + +namespace { + +constexpr size_t kPageSize = 4096; +constexpr size_t kPageCount = 256; + +uint8_t page_value(size_t page, uint8_t bias = 0) { + return static_cast((page * 37 + 11 + bias) & 0xff); +} + +class TemporaryFile { + public: + explicit TemporaryFile(bool unicode_path = false) { + wchar_t temp_dir[MAX_PATH]{}; + DWORD length = ::GetTempPathW(MAX_PATH, temp_dir); + if (length == 0 || length >= static_cast(MAX_PATH) || + ::GetTempFileNameW(temp_dir, L"zvr", 0, wide_path_) == 0) { + wide_path_[0] = L'\0'; + return; + } + + if (unicode_path) { + std::wstring unicode_wide_path = wide_path_; + unicode_wide_path += L"_磁盘索引_テスト"; + if (unicode_wide_path.size() >= static_cast(MAX_PATH) || + !::DeleteFileW(wide_path_)) { + return; + } + std::copy(unicode_wide_path.begin(), unicode_wide_path.end(), wide_path_); + wide_path_[unicode_wide_path.size()] = L'\0'; + } + path_ = zvec::ailego::FileHelper::WideToUtf8(wide_path_); + } + + ~TemporaryFile() { + if (wide_path_[0] != L'\0') { + ::DeleteFileW(wide_path_); + } + } + + bool valid() const { + return wide_path_[0] != L'\0' && !path_.empty(); + } + + const char *path() const { + return path_.c_str(); + } + + const wchar_t *wide_path() const { + return wide_path_; + } + + bool write_pages(uint8_t bias = 0) const { + if (!valid()) { + return false; + } + + HANDLE file = ::CreateFileW(wide_path_, GENERIC_WRITE, 0, nullptr, + CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr); + if (file == INVALID_HANDLE_VALUE) { + return false; + } + + std::vector contents(kPageSize * kPageCount); + for (size_t page = 0; page < kPageCount; ++page) { + std::memset(contents.data() + page * kPageSize, page_value(page, bias), + kPageSize); + } + + DWORD written = 0; + BOOL result = + ::WriteFile(file, contents.data(), static_cast(contents.size()), + &written, nullptr); + bool success = result && written == static_cast(contents.size()) && + ::FlushFileBuffers(file); + ::CloseHandle(file); + return success; + } + + private: + wchar_t wide_path_[MAX_PATH]{}; + std::string path_; +}; + +struct AlignedFree { + void operator()(uint8_t *buffer) const { + ::_aligned_free(buffer); + } +}; + +using AlignedBuffer = std::unique_ptr; + +AlignedBuffer make_aligned_buffer(size_t size) { + auto *buffer = static_cast(::_aligned_malloc(size, kPageSize)); + if (buffer != nullptr) { + std::memset(buffer, 0, size); + } + return AlignedBuffer(buffer); +} + +bool verify_page(const uint8_t *buffer, size_t page, uint8_t bias = 0) { + return std::all_of(buffer, buffer + kPageSize, [page, bias](uint8_t value) { + return value == page_value(page, bias); + }); +} + +DWORD replace_open_file_atomically(const wchar_t *replacement_path, + const wchar_t *target_path) { + // diskann_file_reader.h targets Windows Vista, whose SDK view predates the + // extended rename declarations. These values and this layout are the Win10 + // FILE_RENAME_INFO_EX ABI used by FileRenameInfoEx. + constexpr auto kFileRenameInfoEx = + static_cast(22); + constexpr DWORD kReplaceIfExists = 0x00000001; + constexpr DWORD kPosixSemantics = 0x00000002; + struct ExtendedFileRenameInfo { + DWORD flags; + HANDLE root_directory; + DWORD file_name_length; + WCHAR file_name[MAX_PATH]; + }; + static_assert(offsetof(ExtendedFileRenameInfo, root_directory) == + offsetof(FILE_RENAME_INFO, RootDirectory)); + static_assert(offsetof(ExtendedFileRenameInfo, file_name_length) == + offsetof(FILE_RENAME_INFO, FileNameLength)); + static_assert(offsetof(ExtendedFileRenameInfo, file_name) == + offsetof(FILE_RENAME_INFO, FileName)); + + HANDLE replacement_handle = ::CreateFileW( + replacement_path, DELETE | SYNCHRONIZE, + FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, nullptr, + OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + if (replacement_handle == INVALID_HANDLE_VALUE) { + return ::GetLastError(); + } + + const size_t target_path_length = std::wcslen(target_path); + if (target_path_length >= MAX_PATH) { + ::CloseHandle(replacement_handle); + return ERROR_FILENAME_EXCED_RANGE; + } + + const size_t target_path_bytes = + target_path_length * sizeof(target_path[0]); + ExtendedFileRenameInfo rename_info{}; + rename_info.flags = kReplaceIfExists | kPosixSemantics; + rename_info.root_directory = nullptr; + rename_info.file_name_length = static_cast(target_path_bytes); + std::memcpy(rename_info.file_name, target_path, target_path_bytes); + + // POSIX rename semantics keep existing handles bound to the old file object + // while new opens of the target path observe the replacement. + const BOOL renamed = ::SetFileInformationByHandle( + replacement_handle, kFileRenameInfoEx, &rename_info, + static_cast(sizeof(rename_info))); + const DWORD error = renamed ? ERROR_SUCCESS : ::GetLastError(); + ::CloseHandle(replacement_handle); + return error; +} + +DWORD issue_misaligned_read(HANDLE handle) { + LARGE_INTEGER offset{}; + if (!::SetFilePointerEx(handle, offset, nullptr, FILE_BEGIN)) { + return ::GetLastError(); + } + uint8_t byte = 0; + DWORD bytes_read = 0; + ::SetLastError(ERROR_SUCCESS); + if (::ReadFile(handle, &byte, 1, &bytes_read, nullptr)) { + return ERROR_SUCCESS; + } + return ::GetLastError(); +} + +class ScopedCurrentDirectory { + public: + ScopedCurrentDirectory() { + const DWORD capacity = ::GetCurrentDirectoryW(0, nullptr); + if (capacity == 0) { + return; + } + original_.resize(capacity, L'\0'); + const DWORD length = + ::GetCurrentDirectoryW(capacity, original_.data()); + if (length == 0 || length >= capacity) { + original_.clear(); + return; + } + original_.resize(length); + } + + ~ScopedCurrentDirectory() { + restore(); + } + + bool valid() const { + return !original_.empty(); + } + + bool change_to(const std::wstring &directory) { + return ::SetCurrentDirectoryW(directory.c_str()) != FALSE; + } + + bool restore() { + if (original_.empty()) { + return false; + } + return ::SetCurrentDirectoryW(original_.c_str()) != FALSE; + } + + private: + std::wstring original_; +}; + +} // namespace + +TEST(DiskAnnFileReaderWindowsTest, OpenKeepsStableHandleUnbuffered) { + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + + WindowsAlignedFileReader reader; + reader.open(file.path()); + + HANDLE stable_handle = + WindowsAlignedFileReaderTestPeer::stable_file_handle(reader); + ASSERT_NE(stable_handle, INVALID_HANDLE_VALUE); + EXPECT_EQ(issue_misaligned_read(stable_handle), ERROR_INVALID_PARAMETER); +} + +TEST(DiskAnnFileReaderWindowsTest, OpenSupportsUtf8Path) { + TemporaryFile file(/*unicode_path=*/true); + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + + // Exercise the UTF-8 to UTF-16 conversion used by open() before CreateFileW. + WindowsAlignedFileReader reader; + reader.open(file.path()); + + HANDLE stable_handle = + WindowsAlignedFileReaderTestPeer::stable_file_handle(reader); + ASSERT_NE(stable_handle, INVALID_HANDLE_VALUE); + + IOContext ctx = nullptr; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_NE(ctx, nullptr); + + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector request{{0, kPageSize, output.get()}}; + EXPECT_EQ(reader.read(request, ctx, false), 0); + EXPECT_TRUE(verify_page(output.get(), 0)); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + EXPECT_EQ(ctx, nullptr); +} + +TEST(DiskAnnFileReaderWindowsTest, + RelativePathSurvivesWorkingDirectoryChange) { + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + + const std::wstring full_path(file.wide_path()); + const size_t separator = full_path.find_last_of(L"\\/"); + ASSERT_NE(separator, std::wstring::npos); + const std::wstring directory = full_path.substr(0, separator); + const std::wstring filename = full_path.substr(separator + 1); + + ScopedCurrentDirectory current_directory; + ASSERT_TRUE(current_directory.valid()); + ASSERT_TRUE(current_directory.change_to(directory)); + + WindowsAlignedFileReader reader; + reader.open(zvec::ailego::FileHelper::WideToUtf8(filename)); + + // prepare_io_ctx() opens the actual IOCP handle. Moving away from the + // directory used by open() must not change which file it resolves. + ASSERT_TRUE(current_directory.restore()); + IOContext ctx = nullptr; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_NE(ctx, nullptr); + + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector request{{0, kPageSize, output.get()}}; + EXPECT_EQ(reader.read(request, ctx, false), 0); + EXPECT_TRUE(verify_page(output.get(), 0)); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + EXPECT_EQ(ctx, nullptr); +} + +TEST(DiskAnnFileReaderWindowsTest, ConcurrentContextsKeepCompletionsIsolated) { + constexpr size_t kThreadCount = 8; + constexpr size_t kReadsPerThread = 64; + + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + + WindowsAlignedFileReader reader; + reader.open(file.path()); + + std::atomic ready{0}; + std::atomic start{false}; + std::atomic failures{0}; + std::vector workers; + workers.reserve(kThreadCount); + + for (size_t thread_index = 0; thread_index < kThreadCount; ++thread_index) { + workers.emplace_back([&, thread_index]() { + bool success = true; + IOContext ctx = nullptr; + if (setup_io_ctx(ctx) != 0 || ctx == nullptr) { + success = false; + } + + AlignedBuffer output = make_aligned_buffer(kPageSize * kReadsPerThread); + if (output == nullptr) { + success = false; + } + + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + + if (success) { + std::vector requests; + std::vector source_pages; + requests.reserve(kReadsPerThread); + source_pages.reserve(kReadsPerThread); + for (size_t i = 0; i < kReadsPerThread; ++i) { + const size_t source_page = (thread_index * 13 + i * 7) % kPageCount; + source_pages.push_back(source_page); + requests.emplace_back(source_page * kPageSize, kPageSize, + output.get() + i * kPageSize); + } + + PendingBatch batch; + if (reader.submit(batch, requests, ctx) != 0 || batch.n_reaped != 0 || + ctx->outstanding_count != batch.n_submitted) { + success = false; + } + + std::vector seen(kReadsPerThread, 0); + while (success && batch.n_reaped < batch.n_submitted) { + std::vector completed; + int count = reader.get_completed(batch, ctx, 1, completed); + if (count <= 0 || static_cast(count) != completed.size()) { + success = false; + break; + } + for (uint32_t index : completed) { + if (index >= seen.size() || seen[index] != 0 || + !verify_page(output.get() + index * kPageSize, + source_pages[index])) { + success = false; + break; + } + seen[index] = 1; + } + } + + if (success && !std::all_of(seen.begin(), seen.end(), + [](uint8_t value) { return value == 1; })) { + success = false; + } + } + + if (destroy_io_ctx(ctx) != 0 || ctx != nullptr) { + success = false; + } + if (!success) { + failures.fetch_add(1, std::memory_order_relaxed); + } + }); + } + + while (ready.load(std::memory_order_acquire) != kThreadCount) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + + for (std::thread &worker : workers) { + worker.join(); + } + EXPECT_EQ(failures.load(), 0U); +} + +TEST(DiskAnnFileReaderWindowsTest, ContextCanMoveBetweenRunnableThreads) { + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + + WindowsAlignedFileReader reader; + reader.open(file.path()); + IOContext ctx = nullptr; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_NE(ctx, nullptr); + + AlignedBuffer first_output = make_aligned_buffer(kPageSize); + AlignedBuffer second_output = make_aligned_buffer(kPageSize); + ASSERT_NE(first_output, nullptr); + ASSERT_NE(second_output, nullptr); + + std::atomic first_completed{false}; + std::atomic release_first{false}; + std::atomic second_completed{false}; + int first_result = IndexError_Runtime; + int second_result = IndexError_Runtime; + + std::thread first([&]() { + std::vector requests{{0, kPageSize, first_output.get()}}; + first_result = reader.read(requests, ctx, false); + first_completed.store(true, std::memory_order_release); + + // Stay runnable after dequeuing from the IOCP. With a port concurrency of + // one this thread occupies the only slot even though its batch is done. + while (!release_first.load(std::memory_order_acquire)) { + YieldProcessor(); + } + }); + + while (!first_completed.load(std::memory_order_acquire)) { + ::Sleep(1); + } + + std::thread second([&]() { + std::vector requests{ + {kPageSize, kPageSize, second_output.get()}}; + second_result = reader.read(requests, ctx, false); + second_completed.store(true, std::memory_order_release); + }); + + constexpr ULONGLONG kHandoffTimeoutMs = 5000; + const ULONGLONG deadline = ::GetTickCount64() + kHandoffTimeoutMs; + while (!second_completed.load(std::memory_order_acquire) && + ::GetTickCount64() < deadline) { + ::Sleep(1); + } + const bool completed_while_first_runnable = + second_completed.load(std::memory_order_acquire); + + // Always let the first thread exit before asserting. The old concurrency=1 + // behavior then releases its IOCP association, allowing the second thread + // to finish instead of leaving a permanently hung regression test. + release_first.store(true, std::memory_order_release); + first.join(); + second.join(); + + EXPECT_TRUE(completed_while_first_runnable); + EXPECT_EQ(first_result, 0); + EXPECT_EQ(second_result, 0); + EXPECT_TRUE(verify_page(first_output.get(), 0)); + EXPECT_TRUE(verify_page(second_output.get(), 1)); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + EXPECT_EQ(ctx, nullptr); +} + +TEST(DiskAnnFileReaderWindowsTest, DestroyContextDrainsOutstandingBatch) { + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + + WindowsAlignedFileReader reader; + reader.open(file.path()); + IOContext ctx = nullptr; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_NE(ctx, nullptr); + + AlignedBuffer output = make_aligned_buffer(kPageSize * MAX_IO_DEPTH); + ASSERT_NE(output, nullptr); + + std::vector requests; + requests.reserve(MAX_IO_DEPTH); + for (size_t i = 0; i < MAX_IO_DEPTH; ++i) { + // A zero-byte ReadFile normally completes synchronously. Because the + // reader does not opt into FILE_SKIP_COMPLETION_PORT_ON_SUCCESS, teardown + // must still dequeue its completion packet along with pending reads. + const size_t length = i == 0 ? 0 : kPageSize; + requests.emplace_back(i * kPageSize, length, + output.get() + i * kPageSize); + } + + PendingBatch batch; + ASSERT_EQ(reader.submit(batch, requests, ctx), 0); + ASSERT_EQ(ctx->outstanding_count, batch.n_submitted); + + // This must cancel and reap the batch before it destroys the OVERLAPPED + // slots or lets the caller release output. + ASSERT_EQ(destroy_io_ctx(ctx), 0); + ASSERT_EQ(ctx, nullptr); + + IOContext replacement_ctx = nullptr; + ASSERT_EQ(setup_io_ctx(replacement_ctx), 0); + ASSERT_NE(replacement_ctx, nullptr); + std::vector one_read{{0, kPageSize, output.get()}}; + ASSERT_EQ(reader.read(one_read, replacement_ctx, false), 0); + EXPECT_TRUE(verify_page(output.get(), 0)); + EXPECT_EQ(destroy_io_ctx(replacement_ctx), 0); + EXPECT_EQ(replacement_ctx, nullptr); +} + +TEST(DiskAnnFileReaderWindowsTest, RejectsMisalignedUnbufferedRead) { + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + + WindowsAlignedFileReader reader; + reader.open(file.path()); + IOContext ctx = nullptr; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_NE(ctx, nullptr); + + AlignedBuffer output = make_aligned_buffer(kPageSize * 2); + ASSERT_NE(output, nullptr); + std::vector requests{{0, kPageSize, output.get() + 1}}; + + PendingBatch batch; + EXPECT_EQ(reader.submit(batch, requests, ctx), IndexError_InvalidArgument); + EXPECT_EQ(ctx->outstanding_count, 0U); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + EXPECT_EQ(ctx, nullptr); +} + +TEST(DiskAnnFileReaderWindowsTest, + ReleaseCompletedContextDropsFileHandleAndCanReadAgain) { + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + + WindowsAlignedFileReader reader; + reader.open(file.path()); + IOContext ctx = nullptr; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_NE(ctx, nullptr); + + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector request{{0, kPageSize, output.get()}}; + ASSERT_EQ(reader.read(request, ctx, false), 0); + EXPECT_NE(ctx->file_handle, INVALID_HANDLE_VALUE); + EXPECT_NE(ctx->completion_port, nullptr); + reader.release_io_ctx(ctx); + EXPECT_EQ(ctx->file_handle, INVALID_HANDLE_VALUE); + EXPECT_EQ(ctx->completion_port, nullptr); + EXPECT_EQ(ctx->outstanding_count, 0U); + + // Search contexts may outlive the reader because the high-level context pool + // retains them. Releasing at the complete operation boundary must drop the + // private context handle. The reader's stable handle still owns the old + // contents and can lazily prepare this same context again after deletion. + ASSERT_TRUE(::DeleteFileW(file.wide_path())); + EXPECT_EQ(reader.read(request, ctx, false), 0); + EXPECT_TRUE(verify_page(output.get(), 0)); + reader.release_io_ctx(ctx); + EXPECT_EQ(ctx->file_handle, INVALID_HANDLE_VALUE); + EXPECT_EQ(ctx->completion_port, nullptr); + EXPECT_EQ(ctx->outstanding_count, 0U); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + EXPECT_EQ(ctx, nullptr); +} + +TEST(DiskAnnFileReaderWindowsTest, + OpenFromHandleSurvivesPathReplacementBeforeHandoff) { + constexpr uint8_t kReplacementBias = 97; + + TemporaryFile original; + TemporaryFile replacement; + ASSERT_TRUE(original.valid()); + ASSERT_TRUE(replacement.valid()); + ASSERT_TRUE(original.write_pages()); + ASSERT_TRUE(replacement.write_pages(kReplacementBias)); + + // This buffered handle represents FileReadStorage, which has already + // supplied metadata from the original file. + zvec::ailego::File source; + ASSERT_TRUE(source.open(original.path(), true, false)); + + ASSERT_EQ(replace_open_file_atomically(replacement.wide_path(), + original.wide_path()), + ERROR_SUCCESS); + + WindowsAlignedFileReader original_reader; + ASSERT_EQ(original_reader.open_from_handle(original.path(), + source.native_handle()), + 0); + HANDLE stable_handle = + WindowsAlignedFileReaderTestPeer::stable_file_handle(original_reader); + ASSERT_NE(stable_handle, INVALID_HANDLE_VALUE); + EXPECT_EQ(issue_misaligned_read(stable_handle), ERROR_INVALID_PARAMETER); + source.close(); + + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + IOContext ctx = nullptr; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_NE(ctx, nullptr); + + std::vector request{{0, kPageSize, output.get()}}; + ASSERT_EQ(original_reader.read(request, ctx, false), 0); + EXPECT_TRUE(verify_page(output.get(), 0)); + + WindowsAlignedFileReader replacement_reader; + replacement_reader.open(original.path()); + ASSERT_EQ(replacement_reader.read(request, ctx, false), 0); + EXPECT_TRUE(verify_page(output.get(), 0, kReplacementBias)); + + EXPECT_EQ(destroy_io_ctx(ctx), 0); + EXPECT_EQ(ctx, nullptr); +} + +TEST(DiskAnnFileReaderWindowsTest, + ReusedContextTracksFileObjectAfterAtomicReplacement) { + constexpr uint8_t kReplacementBias = 83; + + TemporaryFile original; + TemporaryFile replacement; + ASSERT_TRUE(original.valid()); + ASSERT_TRUE(replacement.valid()); + ASSERT_TRUE(original.write_pages()); + ASSERT_TRUE(replacement.write_pages(kReplacementBias)); + + WindowsAlignedFileReader original_reader; + original_reader.open(original.path()); + IOContext shared_ctx = nullptr; + ASSERT_EQ(setup_io_ctx(shared_ctx), 0); + ASSERT_NE(shared_ctx, nullptr); + + // The context has not performed I/O yet. Replacing the path must not make + // its first lazy read observe bytes from a different file object. + ASSERT_EQ(replace_open_file_atomically(replacement.wide_path(), + original.wide_path()), + ERROR_SUCCESS); + + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector request{{0, kPageSize, output.get()}}; + ASSERT_EQ(original_reader.read(request, shared_ctx, false), 0); + EXPECT_TRUE(verify_page(output.get(), 0)); + + WindowsAlignedFileReader replacement_reader; + replacement_reader.open(original.path()); + // Reuse the context that is currently bound to original_reader. The path is + // unchanged, so the readers' file identities must force an IOCP rebind. + ASSERT_EQ(replacement_reader.read(request, shared_ctx, false), 0); + EXPECT_TRUE(verify_page(output.get(), 0, kReplacementBias)); + + // Switching back must likewise restore the original file object rather than + // reuse the replacement reader's handle solely because the paths match. + ASSERT_EQ(original_reader.read(request, shared_ctx, false), 0); + EXPECT_TRUE(verify_page(output.get(), 0)); + + EXPECT_EQ(destroy_io_ctx(shared_ctx), 0); + EXPECT_EQ(shared_ctx, nullptr); +} + +TEST(DiskAnnFileReaderWindowsTest, ShortReadResetsContextForNextBatch) { + TemporaryFile file; + ASSERT_TRUE(file.valid()); + ASSERT_TRUE(file.write_pages()); + + WindowsAlignedFileReader reader; + reader.open(file.path()); + IOContext ctx = nullptr; + ASSERT_EQ(setup_io_ctx(ctx), 0); + ASSERT_NE(ctx, nullptr); + + AlignedBuffer output = make_aligned_buffer(kPageSize); + ASSERT_NE(output, nullptr); + std::vector short_request{ + {kPageCount * kPageSize, kPageSize, output.get()}}; + PendingBatch short_batch; + int result = reader.submit(short_batch, short_request, ctx); + if (result == 0) { + std::vector completed; + result = reader.get_completed(short_batch, ctx, 1, completed); + } + EXPECT_NE(result, 0); + + std::vector valid_request{{0, kPageSize, output.get()}}; + EXPECT_EQ(reader.read(valid_request, ctx, false), 0); + EXPECT_TRUE(verify_page(output.get(), 0)); + EXPECT_EQ(destroy_io_ctx(ctx), 0); + EXPECT_EQ(ctx, nullptr); +} 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..f0f7cf1d8 --- /dev/null +++ b/tests/core/algorithm/diskann/diskann_mobile_compat_test.cc @@ -0,0 +1,778 @@ +// 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. + +#if defined(_WIN32) || defined(_WIN64) +#include +#include +#include +#else +#include +#include +#endif +#include +#include +#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_searcher_entity.h" +#include "diskann_util.h" + +namespace zvec::core { +namespace { + +class TemporaryFile { + public: + TemporaryFile() { +#if defined(_WIN32) || defined(_WIN64) + char temp_directory[MAX_PATH]{}; + char temp_file[MAX_PATH]{}; + if (::GetTempPathA(MAX_PATH, temp_directory) != 0 && + ::GetTempFileNameA(temp_directory, "zvc", 0, temp_file) != 0) { + path_ = temp_file; + fd_ = ::_open(path_.c_str(), _O_BINARY | _O_RDWR); + } +#else + char path[] = "DiskAnnMobileCompatTest.XXXXXX"; + fd_ = ::mkstemp(path); + path_ = path; +#endif + } + + ~TemporaryFile() { + if (fd_ >= 0) { + close_descriptor(fd_); + } + remove_file(path_.c_str()); + } + + TemporaryFile(const TemporaryFile &) = delete; + TemporaryFile &operator=(const TemporaryFile &) = delete; + + int fd() const { + return fd_; + } + + const char *path() const { + return path_.c_str(); + } + + void close() { + if (fd_ >= 0) { + close_descriptor(fd_); + fd_ = -1; + } + } + + void release_descriptor_and_unlink() { + close(); + remove_file(path_.c_str()); + } + + private: + static int close_descriptor(int fd) { +#if defined(_WIN32) || defined(_WIN64) + return ::_close(fd); +#else + return ::close(fd); +#endif + } + + static int remove_file(const char *path) { +#if defined(_WIN32) || defined(_WIN64) + return ::_unlink(path); +#else + return ::unlink(path); +#endif + } + + std::string path_; + int fd_{-1}; +}; + +int64_t WriteAt(int fd, const void *data, size_t length, uint64_t offset) { +#if defined(_WIN32) || defined(_WIN64) + if (::_lseeki64(fd, static_cast<__int64>(offset), SEEK_SET) < 0 || + length > + static_cast((std::numeric_limits::max)())) { + return -1; + } + return ::_write(fd, data, static_cast(length)); +#else + return ::pwrite(fd, data, length, static_cast(offset)); +#endif +} + +class VectorSegment final : public IndexStorage::Segment { + public: + explicit VectorSegment(std::vector data) : data_(std::move(data)) {} + + size_t data_size() const override { + return data_.size(); + } + + uint32_t data_crc() const override { + return 0; + } + + size_t padding_size() const override { + return 0; + } + + size_t capacity() const override { + return data_.size(); + } + + size_t fetch(size_t offset, void *buffer, size_t length) const override { + if (offset > data_.size()) { + return 0; + } + const size_t read_size = std::min(length, data_.size() - offset); + if (read_size != 0) { + std::memcpy(buffer, data_.data() + offset, read_size); + } + return read_size; + } + + size_t read(size_t offset, const void **data, size_t length) override { + if (offset > data_.size()) { + *data = nullptr; + return 0; + } + const size_t read_size = std::min(length, data_.size() - offset); + *data = read_size == 0 ? nullptr : data_.data() + offset; + return read_size; + } + + size_t read(size_t offset, IndexStorage::MemoryBlock &data, + size_t length) override { + const void *read_data = nullptr; + const size_t read_size = read(offset, &read_data, length); + data.reset(const_cast(read_data)); + return read_size; + } + + size_t write(size_t, const void *, size_t) override { + return 0; + } + + size_t resize(size_t) override { + return 0; + } + + void update_data_crc(uint32_t) override {} + + Pointer clone() override { + return std::make_shared(data_); + } + + private: + std::vector data_; +}; + +class VectorStorage final : public IndexStorage { + public: + void add(const std::string &id, std::vector data) { + segments_[id] = std::make_shared(std::move(data)); + } + + int init(const ailego::Params &) override { + return 0; + } + + int cleanup() override { + return 0; + } + + int open(const std::string &, bool) override { + return 0; + } + + int flush() override { + return 0; + } + + int close() override { + return 0; + } + + int append(const std::string &, size_t) override { + return IndexError_NotImplemented; + } + + void refresh(uint64_t) override {} + + uint64_t check_point() const override { + return 0; + } + + Segment::Pointer get(const std::string &id, int = -1) override { + const auto it = segments_.find(id); + return it == segments_.end() ? nullptr : it->second; + } + + bool has(const std::string &id) const override { + return segments_.find(id) != segments_.end(); + } + + uint32_t magic() const override { + return 0; + } + + private: + std::map segments_; +}; + +template +void AppendBytes(std::vector *bytes, const T &value) { + const auto *begin = reinterpret_cast(&value); + bytes->insert(bytes->end(), begin, begin + sizeof(value)); +} + +template +std::vector ToBytes(const std::vector &values) { + std::vector bytes(values.size() * sizeof(T)); + if (!bytes.empty()) { + std::memcpy(bytes.data(), values.data(), bytes.size()); + } + return bytes; +} + +std::shared_ptr MakeMinimalEntityStorage( + const std::vector &chunk_offsets, + const std::vector &key_mapping, + const std::vector &keys = {42}, + uint64_t declared_pivot_size = 2 * sizeof(float) * + PQTable::kPQCentroidNum) { + constexpr uint32_t kDimension = 2; + constexpr uint64_t kChunkCount = 2; + const uint64_t document_count = keys.size(); + + auto storage = std::make_shared(); + + DiskAnnMetaHeader header; + header.doc_cnt = document_count; + header.ndims = kDimension; + storage->add(DiskAnnEntity::kDiskAnnMetaSegmentId, + ToBytes(std::vector{header})); + + DiskAnnPqMeta pq_meta; + pq_meta.full_pivot_data_size = declared_pivot_size; + pq_meta.centroid_data_size = kDimension * sizeof(float); + // Legacy indexes leave this field at zero, so the valid case deliberately + // exercises that compatibility path. + pq_meta.chunk_offsets_size = 0; + pq_meta.chunk_num = kChunkCount; + std::vector pq_meta_data; + AppendBytes(&pq_meta_data, pq_meta); + pq_meta_data.resize(pq_meta_data.size() + + kDimension * sizeof(float) * PQTable::kPQCentroidNum + + kDimension * sizeof(float), + 0); + const std::vector chunk_offset_bytes = ToBytes(chunk_offsets); + pq_meta_data.insert(pq_meta_data.end(), chunk_offset_bytes.begin(), + chunk_offset_bytes.end()); + storage->add(DiskAnnEntity::kDiskAnnPqMetaSegmentId, std::move(pq_meta_data)); + storage->add(DiskAnnEntity::kDiskAnnPqDataSegmentId, + std::vector(document_count * kChunkCount, 0)); + + storage->add(DiskAnnEntity::kDiskAnnKeySegmentId, ToBytes(keys)); + storage->add(DiskAnnEntity::kDiskAnnKeyMappingSegmentId, + ToBytes(key_mapping)); + storage->add(DiskAnnEntity::kDiskAnnEntryPointSegmentId, + ToBytes(std::vector{0})); + storage->add(DiskAnnEntity::kDiskAnnVectorSegmentId, {}); + return storage; +} + +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, MinimalEntityUsesValidatedTypedKeyStorage) { + IndexMeta meta(IndexMeta::DataType::DT_FP32, 2); + // Two uint64_t keys and two uint32_t mapping entries both fit in libc++'s + // small-string storage. This specifically exercises the layout that was + // unsafe when these buffers were strings cast to typed pointers. + auto storage = MakeMinimalEntityStorage({0, 1, 2}, {1, 0}, {84, 42}); + + DiskAnnSearcherEntity entity; + ASSERT_EQ(entity.load(meta, storage), 0); + EXPECT_EQ(entity.get_id(42), 1u); + EXPECT_EQ(entity.get_id(84), 0u); + EXPECT_EQ(entity.get_id(41), kInvalidId); + EXPECT_EQ(entity.get_key(0), 84u); + EXPECT_EQ(entity.get_key(1), 42u); + EXPECT_EQ(entity.get_key(2), kInvalidKey); + + const auto cloned = entity.clone(); + ASSERT_NE(cloned, nullptr); + EXPECT_EQ(cloned->get_id(42), 1u); + EXPECT_EQ(cloned->get_key(0), 84u); + const auto cloned_entity = + std::dynamic_pointer_cast(cloned); + ASSERT_NE(cloned_entity, nullptr); + EXPECT_EQ(&entity.entrypoints(), &cloned_entity->entrypoints()); +} + +TEST(DiskAnnMobileCompatTest, EntityAllowsMultipleInvalidKeySlots) { + IndexMeta meta(IndexMeta::DataType::DT_FP32, 2); + auto storage = MakeMinimalEntityStorage({0, 1, 2}, {0, 1, 2}, + {42, kInvalidKey, kInvalidKey}); + + DiskAnnSearcherEntity entity; + ASSERT_EQ(entity.load(meta, storage), 0); + EXPECT_EQ(entity.get_id(42), 0u); + EXPECT_EQ(entity.get_id(kInvalidKey), kInvalidId); + EXPECT_EQ(entity.get_key(1), kInvalidKey); + EXPECT_EQ(entity.get_key(2), kInvalidKey); +} + +TEST(DiskAnnMobileCompatTest, EntityRejectsMalformedMetadata) { + IndexMeta meta(IndexMeta::DataType::DT_FP32, 2); + + { + DiskAnnSearcherEntity entity; + auto storage = MakeMinimalEntityStorage( + {0, 1, 2}, {0}, {42}, 2 * sizeof(float) * PQTable::kPQCentroidNum - 1); + EXPECT_EQ(entity.load(meta, storage), IndexError_InvalidFormat); + } + + { + DiskAnnSearcherEntity entity; + auto storage = MakeMinimalEntityStorage( + {0, 1, 2}, {0}, {42}, std::numeric_limits::max()); + EXPECT_EQ(entity.load(meta, storage), IndexError_InvalidFormat); + } + + { + DiskAnnSearcherEntity entity; + auto storage = MakeMinimalEntityStorage({0, 2, 2}, {0}); + EXPECT_EQ(entity.load(meta, storage), IndexError_InvalidFormat); + } + + { + DiskAnnSearcherEntity entity; + auto storage = MakeMinimalEntityStorage( + {0, 1, 2}, {std::numeric_limits::max()}); + EXPECT_EQ(entity.load(meta, storage), IndexError_InvalidFormat); + } + + { + DiskAnnSearcherEntity entity; + auto storage = MakeMinimalEntityStorage({0, 1, 2}, {0, 0}, {42, 84}); + EXPECT_EQ(entity.load(meta, storage), IndexError_InvalidFormat); + } + + { + DiskAnnSearcherEntity entity; + auto storage = MakeMinimalEntityStorage({0, 1, 2}, {0}); + DiskAnnMetaHeader header; + header.doc_cnt = 1; + header.ndims = 3; + storage->add(DiskAnnEntity::kDiskAnnMetaSegmentId, + ToBytes(std::vector{header})); + EXPECT_EQ(entity.load(meta, storage), IndexError_InvalidFormat); + } + + { + DiskAnnSearcherEntity entity; + auto storage = MakeMinimalEntityStorage({0, 1, 2}, {0}); + storage->add(DiskAnnEntity::kDiskAnnEntryPointSegmentId, + ToBytes(std::vector{1, 1})); + EXPECT_EQ(entity.load(meta, storage), IndexError_InvalidFormat); + } + + { + DiskAnnSearcherEntity entity; + auto storage = MakeMinimalEntityStorage({0, 1, 2}, {0}); + storage->add(DiskAnnEntity::kDiskAnnEntryPointSegmentId, + ToBytes(std::vector{2, 0, 0})); + EXPECT_EQ(entity.load(meta, storage), IndexError_InvalidFormat); + } + + { + DiskAnnSearcherEntity entity; + auto storage = MakeMinimalEntityStorage({0, 1, 2}, {1, 0}, {42, 84}); + EXPECT_EQ(entity.load(meta, storage), IndexError_InvalidFormat); + } +} + +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(WriteAt(file.fd(), expected.data(), expected.size(), 0), + static_cast(expected.size())); + file.close(); + + void *output = nullptr; + DiskAnnUtil::alloc_aligned(&output, kDataSize, kBlockSize); + ASSERT_NE(output, nullptr); + std::memset(output, 0, kDataSize); + + PlatformAlignedFileReader 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(WriteAt(file.fd(), expected.data(), expected.size(), 0), + static_cast(expected.size())); + file.close(); + + void *output = nullptr; + DiskAnnUtil::alloc_aligned(&output, kBlockSize * 2, kBlockSize); + ASSERT_NE(output, nullptr); + + PlatformAlignedFileReader 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(WriteAt(file.fd(), expected.data(), expected.size(), 0), + static_cast(expected.size())); + file.close(); + + void *output = nullptr; + DiskAnnUtil::alloc_aligned(&output, kBlockSize, kBlockSize); + ASSERT_NE(output, nullptr); + + PlatformAlignedFileReader 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(WriteAt(file.fd(), expected.data(), expected.size(), 0), + static_cast(expected.size())); + file.close(); + + std::array outputs{}; + for (void *&output : outputs) { + DiskAnnUtil::alloc_aligned(&output, kBlockSize, kBlockSize); + ASSERT_NE(output, nullptr); + } + + PlatformAlignedFileReader 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); + + std::ifstream snapshot_input(index_file.path(), + std::ios::binary | std::ios::ate); + ASSERT_TRUE(snapshot_input.is_open()); + const std::streamsize snapshot_size = snapshot_input.tellg(); + ASSERT_GT(snapshot_size, 4096); + snapshot_input.seekg(0); + std::vector snapshot(static_cast(snapshot_size)); + ASSERT_TRUE(snapshot_input.read(reinterpret_cast(snapshot.data()), + snapshot_size)); + snapshot_input.close(); + + 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); + + switching_context.reset(); + ASSERT_EQ(first_streamer->close(), 0); + ASSERT_EQ(second_streamer->close(), 0); + first_streamer.reset(); + second_streamer.reset(); + context.reset(); + searcher.reset(); + storage.reset(); + + ASSERT_NO_THROW( + std::filesystem::resize_file(index_file.path(), snapshot.size() - 4096)); + 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(); + std::ofstream restore_output(index_file.path(), + std::ios::binary | std::ios::trunc); + ASSERT_TRUE(restore_output.is_open()); + restore_output.write(reinterpret_cast(snapshot.data()), + static_cast(snapshot.size())); + restore_output.flush(); + ASSERT_TRUE(restore_output.good()); + restore_output.close(); + + 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/core/algorithm/diskann/diskann_node_layout_test.cc b/tests/core/algorithm/diskann/diskann_node_layout_test.cc new file mode 100644 index 000000000..ce9e67c47 --- /dev/null +++ b/tests/core/algorithm/diskann/diskann_node_layout_test.cc @@ -0,0 +1,289 @@ +// 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 "diskann_params.h" +#include "diskann_searcher.h" +#include "diskann_util.h" + +namespace zvec { +namespace core { + +struct DiskAnnNodeLayoutForTest { + uint64_t index_segment_offset{0}; + uint32_t node_per_sector{0}; + uint32_t max_node_size{0}; + uint32_t max_degree{0}; + uint64_t doc_count{0}; + diskann_id_t medoid{0}; +}; + +class DiskAnnCacheTestPeer { + public: + static void set_parser_bounds(DiskAnnIndexer &indexer, uint32_t max_degree, + uint64_t doc_count) { + indexer.max_degree_ = max_degree; + indexer.doc_cnt_ = doc_count; + } + + static int parse_node_neighbors(DiskAnnIndexer &indexer, + const uint8_t *node_buf, diskann_id_t node_id, + uint32_t &neighbor_count, + diskann_id_t *neighbors) { + return indexer.parse_node_neighbors(node_buf, node_id, neighbor_count, + neighbors); + } + + static DiskAnnNodeLayoutForTest layout(const DiskAnnSearcher &searcher) { + const DiskAnnIndexer &indexer = *searcher.diskann_indexer_; + return {indexer.index_segment_offset_, + indexer.node_per_sector_, + indexer.max_node_size_, + indexer.max_degree_, + indexer.doc_cnt_, + indexer.medoid_}; + } +}; + +} // namespace core +} // namespace zvec + +namespace { + +using namespace zvec::ailego; +using namespace zvec::core; + +void ExpectUnalignedNeighborParse(IndexMeta::DataType data_type, + uint32_t dimension) { + IndexMeta meta(data_type, dimension); + DiskAnnIndexer indexer(meta); + DiskAnnCacheTestPeer::set_parser_bounds(indexer, 2, 4); + + std::vector storage(meta.element_size() + sizeof(uint32_t) + + 2 * sizeof(diskann_id_t) + 4); + size_t prefix = 0; + while ((reinterpret_cast(storage.data() + prefix + + meta.element_size()) % + alignof(uint32_t)) == 0) { + ++prefix; + } + uint8_t *node = storage.data() + prefix; + EXPECT_NE(reinterpret_cast(node + meta.element_size()) % + alignof(uint32_t), + 0U); + + const uint32_t stored_count = 2; + const std::array stored_neighbors{1, 3}; + memcpy(node + meta.element_size(), &stored_count, sizeof(stored_count)); + memcpy(node + meta.element_size() + sizeof(stored_count), + stored_neighbors.data(), + stored_neighbors.size() * sizeof(diskann_id_t)); + + uint32_t parsed_count = 0; + std::array parsed_neighbors{}; + EXPECT_EQ(0, DiskAnnCacheTestPeer::parse_node_neighbors( + indexer, node, 0, parsed_count, parsed_neighbors.data())); + EXPECT_EQ(parsed_count, stored_count); + EXPECT_EQ(parsed_neighbors, stored_neighbors); +} + +#if !defined(_WIN32) && !defined(_WIN64) +bool ReadExact(const std::string &path, uint64_t offset, void *data, + size_t size) { + std::ifstream stream(path, std::ios::binary); + if (!stream) { + return false; + } + stream.seekg(static_cast(offset)); + stream.read(static_cast(data), static_cast(size)); + return static_cast(stream.gcount()) == size; +} + +bool WriteExact(const std::string &path, uint64_t offset, const void *data, + size_t size) { + std::fstream stream(path, std::ios::binary | std::ios::in | std::ios::out); + if (!stream) { + return false; + } + stream.seekp(static_cast(offset)); + stream.write(static_cast(data), + static_cast(size)); + stream.flush(); + return stream.good(); +} +#endif + +class DiskAnnNodeLayoutTest : public testing::Test { + protected: + void TearDown() override { + std::error_code error; + std::filesystem::remove(index_path_, error); + } + + const std::string index_path_{"DiskAnnNodeLayoutTest.index"}; +}; + +TEST(DiskAnnNodeParserTest, SafelyParsesUnalignedFp16AndInt8Records) { + ExpectUnalignedNeighborParse(IndexMeta::DataType::DT_FP16, 3); + ExpectUnalignedNeighborParse(IndexMeta::DataType::DT_INT8, 3); +} + +TEST(DiskAnnNodeParserTest, RejectsMalformedCountAndNeighborId) { + IndexMeta meta(IndexMeta::DataType::DT_FP16, 3); + DiskAnnIndexer indexer(meta); + DiskAnnCacheTestPeer::set_parser_bounds(indexer, 2, 4); + + std::vector storage(1 + meta.element_size() + sizeof(uint32_t) + + 3 * sizeof(diskann_id_t)); + uint8_t *node = storage.data() + 1; + std::array output{}; + uint32_t parsed_count = 0; + + uint32_t stored_count = 3; + memcpy(node + meta.element_size(), &stored_count, sizeof(stored_count)); + EXPECT_EQ(IndexError_InvalidFormat, + DiskAnnCacheTestPeer::parse_node_neighbors( + indexer, node, 0, parsed_count, output.data())); + + stored_count = 2; + const std::array stored_neighbors{1, 4}; + memcpy(node + meta.element_size(), &stored_count, sizeof(stored_count)); + memcpy(node + meta.element_size() + sizeof(stored_count), + stored_neighbors.data(), + stored_neighbors.size() * sizeof(diskann_id_t)); + EXPECT_EQ(IndexError_InvalidFormat, + DiskAnnCacheTestPeer::parse_node_neighbors( + indexer, node, 0, parsed_count, output.data())); +} + +TEST_F(DiskAnnNodeLayoutTest, + OddDimensionFp16SearchRejectsMalformedRealtimeNeighbors) { + constexpr uint32_t kDimension = 3; + constexpr size_t kDocCount = 64; + + IndexMeta meta(IndexMeta::DataType::DT_FP16, kDimension); + meta.set_metric("SquaredEuclidean", 0, Params()); + ASSERT_NE(meta.element_size() % alignof(uint32_t), 0U); + + auto holder = + std::make_shared>( + kDimension); + for (size_t i = 0; i < kDocCount; ++i) { + NumericalVector vector(kDimension); + for (size_t d = 0; d < kDimension; ++d) { + vector[d] = static_cast(i + d) / 10.0f; + } + ASSERT_TRUE(holder->emplace(i, vector)); + } + + Params build_params; + build_params.set(PARAM_DISKANN_BUILDER_MAX_DEGREE, 16); + build_params.set(PARAM_DISKANN_BUILDER_LIST_SIZE, 32); + build_params.set(PARAM_DISKANN_BUILDER_MAX_PQ_CHUNK_NUM, 1); + build_params.set(PARAM_DISKANN_BUILDER_THREAD_COUNT, 2); + + auto builder = IndexFactory::CreateBuilder("DiskAnnBuilder"); + ASSERT_NE(builder, nullptr); + ASSERT_EQ(0, builder->init(meta, build_params)); + ASSERT_EQ(0, builder->train(holder)); + ASSERT_EQ(0, builder->build(holder)); + + auto dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_NE(dumper, nullptr); + ASSERT_EQ(0, dumper->create(index_path_)); + ASSERT_EQ(0, builder->dump(dumper)); + ASSERT_EQ(0, dumper->close()); + + Params search_params; + search_params.set(PARAM_DISKANN_SEARCHER_LIST_SIZE, 32); + search_params.set(PARAM_DISKANN_SEARCHER_CACHE_NODE_NUM, 0); + auto searcher = IndexFactory::CreateSearcher("DiskAnnSearcher"); + ASSERT_NE(searcher, nullptr); + ASSERT_EQ(0, searcher->init(search_params)); + + auto storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(storage, nullptr); + ASSERT_EQ(0, storage->open(index_path_, false)); + ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); + + NumericalVector query(kDimension, static_cast(12.1f)); + IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP16, kDimension); + auto run_search = [&]() -> int { + auto context = searcher->create_context(); + if (!context) { + return IndexError_NoMemory; + } + context->set_topk(5); + return searcher->search_impl(query.data(), query_meta, context); + }; + + ASSERT_EQ(0, run_search()); + +#if !defined(_WIN32) && !defined(_WIN64) + auto *diskann_searcher = dynamic_cast(searcher.get()); + ASSERT_NE(diskann_searcher, nullptr); + const DiskAnnNodeLayoutForTest layout = + DiskAnnCacheTestPeer::layout(*diskann_searcher); + ASSERT_GT(layout.max_degree, 0U); + ASSERT_EQ(layout.doc_count, kDocCount); + + const uint64_t node_sector = + DiskAnnUtil::get_node_sector(layout.node_per_sector, layout.max_node_size, + DiskAnnUtil::kSectorSize, layout.medoid); + const uint64_t node_offset = + layout.node_per_sector == 0 + ? 0 + : static_cast(layout.medoid % layout.node_per_sector) * + layout.max_node_size; + const uint64_t count_offset = layout.index_segment_offset + + node_sector * DiskAnnUtil::kSectorSize + + node_offset + meta.element_size(); + + uint32_t original_count = 0; + diskann_id_t original_neighbor = 0; + ASSERT_TRUE(ReadExact(index_path_, count_offset, &original_count, + sizeof(original_count))); + ASSERT_GT(original_count, 0U); + ASSERT_TRUE(ReadExact(index_path_, count_offset + sizeof(original_count), + &original_neighbor, sizeof(original_neighbor))); + + const uint32_t invalid_count = layout.max_degree + 1; + ASSERT_TRUE(WriteExact(index_path_, count_offset, &invalid_count, + sizeof(invalid_count))); + EXPECT_EQ(IndexError_InvalidFormat, run_search()); + + ASSERT_TRUE(WriteExact(index_path_, count_offset, &original_count, + sizeof(original_count))); + const diskann_id_t invalid_neighbor = + static_cast(layout.doc_count); + ASSERT_TRUE(WriteExact(index_path_, count_offset + sizeof(original_count), + &invalid_neighbor, sizeof(invalid_neighbor))); + EXPECT_EQ(IndexError_InvalidFormat, run_search()); + + ASSERT_TRUE(WriteExact(index_path_, count_offset + sizeof(original_count), + &original_neighbor, sizeof(original_neighbor))); + EXPECT_EQ(0, run_search()); +#endif +} + +} // namespace diff --git a/tests/core/algorithm/diskann/diskann_searcher_test.cc b/tests/core/algorithm/diskann/diskann_searcher_test.cc index 84bf8eb17..1292797ef 100644 --- a/tests/core/algorithm/diskann/diskann_searcher_test.cc +++ b/tests/core/algorithm/diskann/diskann_searcher_test.cc @@ -13,28 +13,284 @@ // limitations under the License. #include "diskann_searcher.h" -#include -#include -#include -#include +#include #include +#include #include +#include +#include #include #include +#include #include #include #include #include #include -#include "diskann_holder.h" +#include "diskann_index_provider.h" #include "diskann_params.h" +#include "diskann_streamer.h" +#include "diskann_util.h" + +namespace zvec { +namespace core { + +class DiskAnnCacheTestPeer { + public: + static std::shared_ptr reader(DiskAnnSearcher *searcher) { + return searcher->diskann_indexer_->reader_; + } + + static void set_reader(DiskAnnSearcher *searcher, + std::shared_ptr reader) { + searcher->diskann_indexer_->reader_ = std::move(reader); + } + + static int configure_cache(DiskAnnSearcher *searcher, + uint32_t cache_node_num) { + return searcher->diskann_indexer_->configure_cache(cache_node_num); + } + + static size_t coordinate_cache_size(const DiskAnnSearcher *searcher) { + return searcher->diskann_indexer_->coord_cache_.size(); + } + + static size_t neighbor_cache_size(const DiskAnnSearcher *searcher) { + return searcher->diskann_indexer_->neighbor_cache_.size(); + } + + static diskann_key_t replace_key(DiskAnnSearcher *searcher, diskann_id_t id, + diskann_key_t replacement) { + auto *entity = dynamic_cast( + searcher->diskann_indexer_->entity_.get()); + if (entity == nullptr || !entity->key_buffer_ || + !entity->key_mapping_buffer_ || + id >= entity->key_mapping_buffer_->size()) { + return kInvalidKey; + } + const diskann_id_t key_index = (*entity->key_mapping_buffer_)[id]; + if (key_index >= entity->key_buffer_->size()) { + return kInvalidKey; + } + auto keys = std::const_pointer_cast>( + entity->key_buffer_); + const diskann_key_t previous = (*keys)[key_index]; + (*keys)[key_index] = replacement; + return previous; + } +}; + +class DiskAnnStreamerTestPeer { + public: + static std::shared_ptr reader(DiskAnnStreamer *streamer) { + return streamer->diskann_indexer_->reader_; + } + + static void set_reader(DiskAnnStreamer *streamer, + std::shared_ptr reader) { + streamer->diskann_indexer_->reader_ = std::move(reader); + } +}; + +class DiskAnnProviderTestPeer { + public: + static DiskAnnContext *fetch_context(IndexProvider *provider) { + auto *diskann_provider = dynamic_cast(provider); + return diskann_provider == nullptr + ? nullptr + : dynamic_cast( + diskann_provider->fetch_context_.get()); + } + + static DiskAnnContext *iterator_context(IndexProvider::Iterator *iterator) { + auto *diskann_iterator = + dynamic_cast(iterator); + return diskann_iterator == nullptr ? nullptr + : dynamic_cast( + diskann_iterator->context_.get()); + } +}; + +} // namespace core +} // namespace zvec using namespace zvec::core; using namespace zvec::ailego; using namespace std; +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); + constexpr size_t static dim = 64; +namespace { + +class CountingAlignedFileReader final : public AlignedFileReader { + public: + explicit CountingAlignedFileReader(std::shared_ptr reader) + : reader_(std::move(reader)) {} + + void open(const std::string &fname) override { + reader_->open(fname); + } + + void close() override { + reader_->close(); + } + + int read(std::vector &read_reqs, IOContext &ctx, + bool async = false) override { + requested_reads_ += read_reqs.size(); + return reader_->read(read_reqs, ctx, async); + } + + int submit(PendingBatch &batch, std::vector &read_reqs, + IOContext &ctx) override { + return reader_->submit(batch, read_reqs, ctx); + } + + int get_completed(PendingBatch &batch, IOContext &ctx, int min_completed, + std::vector &completed_indices) override { + return reader_->get_completed(batch, ctx, min_completed, completed_indices); + } + + void release_io_ctx(IOContext &ctx) override { + ++release_count_; + reader_->release_io_ctx(ctx); + } + + size_t requested_reads() const { + return requested_reads_; + } + + size_t release_count() const { + return release_count_; + } + + private: + std::shared_ptr reader_; + size_t requested_reads_{0}; + size_t release_count_{0}; +}; + +class ContextTrackingAlignedFileReader final : public AlignedFileReader { + public: + explicit ContextTrackingAlignedFileReader( + std::shared_ptr reader) + : reader_(std::move(reader)) {} + + void open(const std::string &fname) override { + reader_->open(fname); + } + + void close() override { + reader_->close(); + } + + int read(std::vector &read_reqs, IOContext &ctx, + bool async = false) override { + { + std::lock_guard lock(mutex_); + contexts_.insert(ctx); + } + return reader_->read(read_reqs, ctx, async); + } + + int submit(PendingBatch &batch, std::vector &read_reqs, + IOContext &ctx) override { + return reader_->submit(batch, read_reqs, ctx); + } + + int get_completed(PendingBatch &batch, IOContext &ctx, int min_completed, + std::vector &completed_indices) override { + return reader_->get_completed(batch, ctx, min_completed, completed_indices); + } + + void release_io_ctx(IOContext &ctx) override { + reader_->release_io_ctx(ctx); + } + + size_t context_count() const { + std::lock_guard lock(mutex_); + return contexts_.size(); + } + + private: + std::shared_ptr reader_; + mutable std::mutex mutex_; + std::set contexts_; +}; + +class FailingAlignedFileReader final : public AlignedFileReader { + public: + explicit FailingAlignedFileReader(std::shared_ptr reader) + : reader_(std::move(reader)) {} + + void open(const std::string &fname) override { + reader_->open(fname); + } + + void close() override { + reader_->close(); + } + + int read(std::vector & /*read_reqs*/, IOContext & /*ctx*/, + bool /*async*/ = false) override { + return IndexError_ReadData; + } + + int submit(PendingBatch & /*batch*/, std::vector & /*read_reqs*/, + IOContext & /*ctx*/) override { + return IndexError_ReadData; + } + + int get_completed(PendingBatch & /*batch*/, IOContext & /*ctx*/, + int /*min_completed*/, + std::vector & /*completed_indices*/) override { + return IndexError_ReadData; + } + + void release_io_ctx(IOContext &ctx) override { + reader_->release_io_ctx(ctx); + } + + private: + std::shared_ptr reader_; +}; + +class CorruptibleDiskAnnSearcherEntity final : public DiskAnnSearcherEntity { + public: + void set_node_layout(uint64_t max_node_size, uint64_t node_per_sector) { + meta_header_.max_node_size = max_node_size; + meta_header_.node_per_sector = node_per_sector; + } + + void set_index_size(uint64_t index_size) { + meta_header_.index_size = index_size; + } +}; + +size_t expected_fetch_buffer_size(const DiskAnnContext &context) { + const auto &entity = context.get_entity(); + const uint64_t sector_num_per_node = + entity.node_per_sector() > 0 + ? 1 + : DiskAnnUtil::div_round_up(entity.max_node_size(), + DiskAnnUtil::kSectorSize); + return static_cast(sector_num_per_node) * DiskAnnUtil::kSectorSize; +} + +void expect_same_results(const IndexDocumentList &expected, + const IndexDocumentList &actual) { + ASSERT_EQ(expected.size(), actual.size()); + for (size_t i = 0; i < expected.size(); ++i) { + EXPECT_EQ(expected[i].key(), actual[i].key()); + EXPECT_FLOAT_EQ(expected[i].score(), actual[i].score()); + } +} + +} // namespace + class DiskAnnSearcherTest : public testing::Test { protected: void SetUp(void) override; @@ -56,9 +312,29 @@ void DiskAnnSearcherTest::SetUp(void) { } void DiskAnnSearcherTest::TearDown(void) { - char cmdBuf[100]; - snprintf(cmdBuf, 100, "rm -rf %s", _dir.c_str()); - system(cmdBuf); + std::filesystem::remove_all(_dir); +} + +TEST_F(DiskAnnSearcherTest, TestRejectZeroSearchListSize) { + Params params; + params.set(PARAM_DISKANN_SEARCHER_LIST_SIZE, 0); + + auto searcher = IndexFactory::CreateSearcher("DiskAnnSearcher"); + ASSERT_NE(searcher, nullptr); + EXPECT_EQ(IndexError_InvalidArgument, searcher->init(params)); + + auto streamer = IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(streamer, nullptr); + EXPECT_EQ(IndexError_InvalidArgument, + streamer->init(*_index_meta_ptr, params)); + + NeighborPriorityQueue default_queue; + default_queue.insert(Neighbor(1, 1.0f)); + EXPECT_EQ(0U, default_queue.size()); + + NeighborPriorityQueue zero_capacity_queue(0); + zero_capacity_queue.insert(Neighbor(1, 1.0f)); + EXPECT_EQ(0U, zero_capacity_queue.size()); } TEST_F(DiskAnnSearcherTest, TestGeneral) { @@ -97,6 +373,47 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) { ASSERT_EQ(0, builder->dump(dumper)); ASSERT_EQ(0, dumper->close()); + // A fetch context now allocates only the sectors required for one node, so + // the on-disk packing fields must be mutually consistent. Otherwise a + // forged nodes-per-sector value can put the computed node offset beyond + // that exact buffer even though the sector read itself fits. + { + auto malformed_storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(malformed_storage, nullptr); + ASSERT_EQ(0, malformed_storage->open(path, false)); + CorruptibleDiskAnnSearcherEntity malformed_entity; + ASSERT_EQ(0, malformed_entity.load(*_index_meta_ptr, malformed_storage)); + const uint64_t max_node_size = malformed_entity.max_node_size(); + const uint64_t index_size = malformed_entity.index_size(); + const uint64_t expected_node_per_sector = + max_node_size <= DiskAnnUtil::kSectorSize + ? DiskAnnUtil::kSectorSize / max_node_size + : 0; + malformed_entity.set_node_layout(max_node_size, + expected_node_per_sector + 1); + DiskAnnIndexer malformed_indexer(*_index_meta_ptr); + EXPECT_EQ(IndexError_InvalidFormat, + malformed_indexer.init(malformed_entity)); + + // Keep the packing formula self-consistent, but remove the space required + // for the declared adjacency list. This must be rejected independently of + // the nodes-per-sector consistency check above. + ASSERT_GT(malformed_entity.max_degree(), 0U); + const uint64_t undersized_node = + _index_meta_ptr->element_size() + sizeof(uint32_t); + malformed_entity.set_node_layout( + undersized_node, DiskAnnUtil::kSectorSize / undersized_node); + DiskAnnIndexer undersized_indexer(*_index_meta_ptr); + EXPECT_EQ(IndexError_InvalidFormat, + undersized_indexer.init(malformed_entity)); + + malformed_entity.set_node_layout(max_node_size, expected_node_per_sector); + malformed_entity.set_index_size(index_size + DiskAnnUtil::kSectorSize); + DiskAnnIndexer wrong_size_indexer(*_index_meta_ptr); + EXPECT_EQ(IndexError_InvalidFormat, + wrong_size_indexer.init(malformed_entity)); + } + auto &stats = builder->stats(); ASSERT_EQ(doc_cnt, stats.trained_count()); ASSERT_EQ(doc_cnt, stats.built_count()); @@ -112,22 +429,89 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) { Params search_params; search_params.set("zvec.diskann.searcher.list_size", 500); + search_params.set("zvec.diskann.searcher.cache_node_num", 0); ASSERT_EQ(0, searcher->init(search_params)); + // Independent FileReadStorage segments do not expose one file object that + // can anchor all DiskAnn segments to the same snapshot. Reject that mode on + // every platform before opening the aligned graph reader. + auto independent_storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(independent_storage, nullptr); + Params independent_storage_params; + independent_storage_params.set("proxima.file.read_storage.alone_file_handle", + true); + ASSERT_EQ(0, independent_storage->init(independent_storage_params)); + ASSERT_EQ(0, independent_storage->open(path, false)); + auto retained_independent_segment = + independent_storage->get(DiskAnnEntity::kDiskAnnVectorSegmentId); + ASSERT_NE(retained_independent_segment, nullptr); + ASSERT_EQ(nullptr, independent_storage->file()); + EXPECT_EQ(IndexError_InvalidArgument, + searcher->load(independent_storage, IndexMetric::Pointer())); + auto independent_streamer = IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(independent_streamer, nullptr); + ASSERT_EQ(0, independent_streamer->init(*_index_meta_ptr, search_params)); + EXPECT_EQ(IndexError_InvalidArgument, + independent_streamer->open(independent_storage)); + { + DiskAnnSearcherEntity independent_entity; + ASSERT_EQ(0, + independent_entity.load(*_index_meta_ptr, independent_storage)); + DiskAnnIndexer independent_indexer(*_index_meta_ptr); + EXPECT_EQ(IndexError_InvalidArgument, + independent_indexer.init(independent_entity)); + } + retained_independent_segment.reset(); + independent_storage.reset(); + auto storage = IndexFactory::CreateStorage("FileReadStorage"); ASSERT_EQ(0, storage->open(path, false)); + auto retained_cached_file = storage->file(); + ASSERT_NE(retained_cached_file, nullptr); + ASSERT_TRUE(retained_cached_file->is_valid()); + auto retained_segment = storage->get(DiskAnnEntity::kDiskAnnVectorSegmentId); + ASSERT_NE(retained_segment, nullptr); + std::weak_ptr searcher_cached_file = retained_cached_file; +#if defined(_WIN32) || defined(_WIN64) + // Keeping an ordinary buffered alias beside DiskAnn's unbuffered handles + // causes a severe random-read regression on Windows. Reject the load without + // invalidating either caller-owned alias. Releasing those aliases allows the + // same, still-open storage to be retried. + EXPECT_EQ(IndexError_InvalidArgument, + searcher->load(storage, IndexMetric::Pointer())); + EXPECT_TRUE(retained_cached_file->is_valid()); + uint8_t retained_file_byte = 0; + EXPECT_EQ(1U, retained_segment->fetch(0, &retained_file_byte, 1)); + retained_cached_file.reset(); + retained_segment.reset(); + EXPECT_FALSE(searcher_cached_file.expired()); ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); + EXPECT_TRUE(searcher_cached_file.expired()); +#else + ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); + // DiskAnn owns an independent descriptor. Loading must not close or enable + // direct I/O on the File shared by caller-owned FileReadStorage segments. + EXPECT_TRUE(retained_cached_file->is_valid()); + uint8_t retained_file_byte = 0; + EXPECT_EQ(1U, retained_segment->fetch(0, &retained_file_byte, 1)); + retained_cached_file.reset(); + EXPECT_FALSE(searcher_cached_file.expired()); + retained_segment.reset(); + EXPECT_TRUE(searcher_cached_file.expired()); +#endif auto ctx = searcher->create_context(); ASSERT_TRUE(!!ctx); auto linearCtx = searcher->create_context(); auto linearByPKeysCtx = searcher->create_context(); auto knnCtx = searcher->create_context(); + auto singleKnnCtx = searcher->create_context(); ASSERT_TRUE(!!linearCtx); ASSERT_TRUE(!!linearByPKeysCtx); ASSERT_TRUE(!!knnCtx); + ASSERT_TRUE(!!singleKnnCtx); NumericalVector vec(dim); IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, dim); @@ -138,6 +522,99 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) { linearCtx->set_topk(topk); linearByPKeysCtx->set_topk(topk); knnCtx->set_topk(topk); + singleKnnCtx->set_topk(topk); + + auto *diskann_knn_ctx = dynamic_cast(knnCtx.get()); + ASSERT_NE(diskann_knn_ctx, nullptr); + ASSERT_EQ(500U, diskann_knn_ctx->list_size()); + Params zero_list_size; + zero_list_size.set(PARAM_DISKANN_SEARCHER_LIST_SIZE, 0); + EXPECT_EQ(IndexError_InvalidArgument, + diskann_knn_ctx->update(zero_list_size)); + EXPECT_EQ(500U, diskann_knn_ctx->list_size()); + + auto *diskann_searcher = dynamic_cast(searcher.get()); + ASSERT_NE(diskann_searcher, nullptr); + + // Deleted or never-populated slots remain graph nodes so traversal can use + // them for connectivity, but they must never reach filters or user-visible + // results. + { + const diskann_key_t original_key = + DiskAnnCacheTestPeer::replace_key(diskann_searcher, 3, kInvalidKey); + ASSERT_EQ(3U, original_key); + + auto invalid_slot_ctx = searcher->create_context(); + ASSERT_NE(invalid_slot_ctx, nullptr); + invalid_slot_ctx->set_topk(8); + bool invalid_key_reached_filter = false; + invalid_slot_ctx->set_filter([&](uint64_t key) { + invalid_key_reached_filter |= key == kInvalidKey; + return false; + }); + + std::array invalid_slot_query{}; + invalid_slot_query.fill(3.1F); + ASSERT_EQ(0, searcher->search_bf_impl(invalid_slot_query.data(), qmeta, + invalid_slot_ctx)); + EXPECT_FALSE(invalid_key_reached_filter); + for (const auto &doc : invalid_slot_ctx->result()) { + EXPECT_NE(kInvalidKey, doc.key()); + EXPECT_NE(3U, doc.key()); + } + + ASSERT_EQ(0, searcher->search_impl(invalid_slot_query.data(), qmeta, + invalid_slot_ctx)); + EXPECT_FALSE(invalid_key_reached_filter); + for (const auto &doc : invalid_slot_ctx->result()) { + EXPECT_NE(kInvalidKey, doc.key()); + EXPECT_NE(3U, doc.key()); + } + + EXPECT_EQ(kInvalidKey, DiskAnnCacheTestPeer::replace_key(diskann_searcher, + 3, original_key)); + } + + auto batch_counting_reader = std::make_shared( + DiskAnnCacheTestPeer::reader(diskann_searcher)); + DiskAnnCacheTestPeer::set_reader(diskann_searcher, batch_counting_reader); + + // A public count>1 call is one I/O lease. Its individual queries may issue + // many batches, but the pooled context must be released exactly once when + // the complete public operation returns. + constexpr uint32_t kBatchQueryCount = 3; + std::vector batch_queries(kBatchQueryCount * dim); + for (uint32_t query_index = 0; query_index < kBatchQueryCount; + ++query_index) { + std::fill(batch_queries.begin() + query_index * dim, + batch_queries.begin() + (query_index + 1) * dim, + static_cast(query_index) + 0.1f); + } + + size_t release_count = batch_counting_reader->release_count(); + ASSERT_EQ(0, searcher->search_impl(batch_queries.data(), qmeta, + kBatchQueryCount, knnCtx)); + EXPECT_EQ(release_count + 1, batch_counting_reader->release_count()); + for (uint32_t query_index = 0; query_index < kBatchQueryCount; + ++query_index) { + SCOPED_TRACE(query_index); + ASSERT_EQ(0, searcher->search_impl(batch_queries.data() + query_index * dim, + qmeta, singleKnnCtx)); + expect_same_results(singleKnnCtx->result(), knnCtx->result(query_index)); + } + + release_count = batch_counting_reader->release_count(); + ASSERT_EQ(0, searcher->search_bf_impl(batch_queries.data(), qmeta, + kBatchQueryCount, linearCtx)); + EXPECT_EQ(release_count + 1, batch_counting_reader->release_count()); + + std::vector> batch_p_keys(kBatchQueryCount, + {0, 1, 2, 3, 4, 5, 6, 7}); + release_count = batch_counting_reader->release_count(); + ASSERT_EQ(0, searcher->search_bf_by_p_keys_impl( + batch_queries.data(), batch_p_keys, qmeta, kBatchQueryCount, + linearByPKeysCtx)); + EXPECT_EQ(release_count + 1, batch_counting_reader->release_count()); // do linear search test { @@ -226,13 +703,40 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) { auto streamer_storage = IndexFactory::CreateStorage("FileReadStorage"); ASSERT_EQ(0, streamer_storage->open(path, false)); + std::weak_ptr streamer_cached_file = + streamer_storage->file(); + ASSERT_FALSE(streamer_cached_file.expired()); ASSERT_EQ(0, streamer->open(streamer_storage)); + EXPECT_TRUE(streamer_cached_file.expired()); + + auto *diskann_streamer = dynamic_cast(streamer.get()); + ASSERT_NE(diskann_streamer, nullptr); + auto streamer_counting_reader = std::make_shared( + DiskAnnStreamerTestPeer::reader(diskann_streamer)); + DiskAnnStreamerTestPeer::set_reader(diskann_streamer, + streamer_counting_reader); auto streamer_ctx = streamer->create_context(); + auto single_streamer_ctx = streamer->create_context(); ASSERT_NE(streamer_ctx, nullptr); + ASSERT_NE(single_streamer_ctx, nullptr); streamer_ctx->set_topk(topk); + single_streamer_ctx->set_topk(topk); auto *original_ctx = streamer_ctx.get(); + release_count = streamer_counting_reader->release_count(); + ASSERT_EQ(0, streamer->search_impl(batch_queries.data(), qmeta, + kBatchQueryCount, streamer_ctx)); + EXPECT_EQ(release_count + 1, streamer_counting_reader->release_count()); + for (uint32_t query_index = 0; query_index < kBatchQueryCount; + ++query_index) { + SCOPED_TRACE(query_index); + ASSERT_EQ(0, streamer->search_impl(batch_queries.data() + query_index * dim, + qmeta, single_streamer_ctx)); + expect_same_results(single_streamer_ctx->result(), + streamer_ctx->result(query_index)); + } + ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); EXPECT_EQ(original_ctx, streamer_ctx.get()); ASSERT_EQ(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); @@ -282,8 +786,13 @@ TEST_F(DiskAnnSearcherTest, TestGeneral) { // I/O failures from the indexer must be propagated by the streamer instead // of being converted into a successful search with incomplete results. - ASSERT_EQ(0, ::truncate(path.c_str(), 0)); + auto failing_reader = std::make_shared( + std::make_shared( + DiskAnnStreamerTestPeer::reader(diskann_streamer))); + DiskAnnStreamerTestPeer::set_reader(diskann_streamer, failing_reader); + release_count = failing_reader->release_count(); EXPECT_NE(0, streamer->search_impl(vec.data(), qmeta, streamer_ctx)); + EXPECT_EQ(release_count + 1, failing_reader->release_count()); // Closing/unloading releases the index and makes all query entry points // reject work until another index is loaded. @@ -347,7 +856,8 @@ TEST_F(DiskAnnSearcherTest, TestNodeCache) { ASSERT_TRUE(searcher != nullptr); Params search_params; - search_params.set("zvec.diskann.searcher.cache_node_num", 32); + constexpr uint32_t kCacheNodes = 2 * DiskAnnUtil::kMaxSectorReadNum + 3; + search_params.set("zvec.diskann.searcher.cache_node_num", kCacheNodes); search_params.set("zvec.diskann.searcher.list_size", 500); ASSERT_EQ(0, searcher->init(search_params)); @@ -355,6 +865,24 @@ TEST_F(DiskAnnSearcherTest, TestNodeCache) { auto storage = IndexFactory::CreateStorage("FileReadStorage"); ASSERT_EQ(0, storage->open(path, false)); ASSERT_EQ(0, searcher->load(storage, IndexMetric::Pointer())); + + // Count all reads made by a second cache build. BFS-expanded nodes are + // written directly into their final cache slots, so every selected node is + // requested from the underlying reader at most once across BFS and the + // final preload pass. + auto *diskann_searcher = dynamic_cast(searcher.get()); + ASSERT_NE(nullptr, diskann_searcher); + auto counting_reader = std::make_shared( + DiskAnnCacheTestPeer::reader(diskann_searcher)); + DiskAnnCacheTestPeer::set_reader(diskann_searcher, counting_reader); + ASSERT_EQ( + 0, DiskAnnCacheTestPeer::configure_cache(diskann_searcher, kCacheNodes)); + EXPECT_EQ(kCacheNodes, + DiskAnnCacheTestPeer::coordinate_cache_size(diskann_searcher)); + EXPECT_EQ(kCacheNodes, + DiskAnnCacheTestPeer::neighbor_cache_size(diskann_searcher)); + EXPECT_EQ(kCacheNodes, counting_reader->requested_reads()); + auto ctx = searcher->create_context(); ASSERT_TRUE(!!ctx); @@ -626,11 +1154,48 @@ TEST_F(DiskAnnSearcherTest, TestGroup) { ctx->set_group_params(group_num, group_topk); ctx->set_group_by(groupbyFunc); - size_t query_value = doc_cnt / 2; + size_t query_value = doc_cnt * 11 / 20; for (size_t j = 0; j < dim; ++j) { vec[j] = query_value / 10 + 0.1f; } + // Force group stage 2 by assigning the regular top two documents to one + // group and every newly discovered document to another. Its PQ query must + // remain identical to a regular search, because stage 1 already centered it + // and built the distance table. + auto regular_ctx = searcher->create_context(); + regular_ctx->set_topk(2); + ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, regular_ctx)); + + std::unordered_set initial_keys; + for (const auto &doc : regular_ctx->result()) { + initial_keys.insert(doc.key()); + } + auto forcedGroupbyFunc = [&initial_keys](uint64_t key) { + return initial_keys.count(key) != 0 ? std::string("initial") + : std::string("expanded"); + }; + auto forced_group_ctx = searcher->create_context(); + forced_group_ctx->set_group_params(2, 1); + forced_group_ctx->set_group_by(forcedGroupbyFunc); + ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, forced_group_ctx)); + + auto *regular_diskann_ctx = dynamic_cast(regular_ctx.get()); + auto *forced_group_diskann_ctx = + dynamic_cast(forced_group_ctx.get()); + ASSERT_NE(nullptr, regular_diskann_ctx); + ASSERT_NE(nullptr, forced_group_diskann_ctx); + EXPECT_EQ(0, std::memcmp(regular_diskann_ctx->query_rotated(), + forced_group_diskann_ctx->query_rotated(), + _index_meta_ptr->element_size())); + + const auto &forced_group_result = forced_group_ctx->group_result(); + ASSERT_EQ(2U, forced_group_result.size()); + for (const auto &group : forced_group_result) { + ASSERT_EQ(1U, group.docs().size()); + EXPECT_EQ(group.group_id(), forcedGroupbyFunc(group.docs()[0].key())); + } + ASSERT_EQ(0, searcher->search_impl(vec.data(), qmeta, ctx)); auto &group_result = ctx->group_result(); @@ -788,6 +1353,12 @@ TEST_F(DiskAnnSearcherTest, TestFetchVector) { auto linearCtx = searcher->create_context(); auto knnCtx = searcher->create_context(); auto linearByPKeysCtx = searcher->create_context(); + auto *diskann_search_context = + dynamic_cast(linearCtx.get()); + ASSERT_NE(diskann_search_context, nullptr); + EXPECT_EQ(static_cast(DiskAnnUtil::kMaxSectorReadNum) * + DiskAnnUtil::kSectorSize, + diskann_search_context->sector_buffer_size()); knnCtx->set_fetch_vector(true); for (size_t i = 0; i < doc_cnt; i += doc_cnt / 10) { @@ -838,6 +1409,160 @@ TEST_F(DiskAnnSearcherTest, TestFetchVector) { searcher->get_vector(42, linearCtx, missing_vector)); EXPECT_TRUE(missing_vector.empty()); + // A DiskAnn provider reads through the aligned index reader rather than a + // FileReadStorage segment. It and its iterator therefore remain usable + // after their source streamer is closed. + IndexStreamer::Pointer streamer = + IndexFactory::CreateStreamer("DiskAnnStreamer"); + ASSERT_NE(streamer, nullptr); + ASSERT_EQ(0, streamer->init(*_index_meta_ptr, search_params)); + auto streamer_storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(streamer_storage, nullptr); + ASSERT_EQ(0, streamer_storage->open(path, false)); + std::weak_ptr provider_cached_file = + streamer_storage->file(); + ASSERT_FALSE(provider_cached_file.expired()); + ASSERT_EQ(0, streamer->open(streamer_storage)); + ASSERT_TRUE(provider_cached_file.expired()); + + auto *diskann_streamer = dynamic_cast(streamer.get()); + ASSERT_NE(diskann_streamer, nullptr); + auto tracking_reader = std::make_shared( + DiskAnnStreamerTestPeer::reader(diskann_streamer)); + DiskAnnStreamerTestPeer::set_reader(diskann_streamer, tracking_reader); + + auto provider = streamer->create_provider(); + ASSERT_NE(provider, nullptr); + auto second_provider = streamer->create_provider(); + ASSERT_NE(second_provider, nullptr); + EXPECT_TRUE(provider_cached_file.expired()); + EXPECT_EQ(doc_cnt, provider->count()); + EXPECT_EQ(dim, provider->dimension()); + EXPECT_EQ(IndexMeta::DataType::DT_FP32, provider->data_type()); + EXPECT_EQ(_index_meta_ptr->element_size(), provider->element_size()); + + float provider_value = 0.0f; + auto provider_iterator = provider->create_iterator(); + ASSERT_NE(provider_iterator, nullptr); + ASSERT_TRUE(provider_iterator->is_valid()); + EXPECT_EQ(key_for_id(0), provider_iterator->key()); + EXPECT_EQ(nullptr, DiskAnnProviderTestPeer::fetch_context(provider.get())); + EXPECT_EQ(nullptr, + DiskAnnProviderTestPeer::iterator_context(provider_iterator.get())); + float iterator_value = 0.0f; + + // Neither object has performed vector I/O yet. Their first lazy context and + // aligned file handle must still be creatable after the streamer closes. + ASSERT_EQ(0, streamer->close()); + + const void *provider_vector = provider->get_vector(key_for_id(17)); + ASSERT_NE(provider_vector, nullptr); + auto *provider_fetch_context = + DiskAnnProviderTestPeer::fetch_context(provider.get()); + ASSERT_NE(provider_fetch_context, nullptr); + EXPECT_EQ(expected_fetch_buffer_size(*provider_fetch_context), + provider_fetch_context->sector_buffer_size()); + EXPECT_LT(provider_fetch_context->sector_buffer_size(), + static_cast(DiskAnnUtil::kMaxSectorReadNum) * + DiskAnnUtil::kSectorSize); + std::memcpy(&provider_value, provider_vector, sizeof(provider_value)); + EXPECT_EQ(17.0f, provider_value); + + // A returned pointer must not be overwritten by a concurrent fetch on a + // different thread. Coordinate the calls so this deterministically catches + // providers that share one result buffer globally. + std::atomic first_fetch_ready{false}; + std::atomic second_fetch_done{false}; + std::atomic release_first_fetch{false}; + float first_thread_value = -1.0f; + float second_thread_value = -1.0f; + std::thread first_fetch([&]() { + const void *value = provider->get_vector(key_for_id(31)); + first_fetch_ready.store(true, std::memory_order_release); + while (!second_fetch_done.load(std::memory_order_acquire) && + !release_first_fetch.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + if (value != nullptr) { + std::memcpy(&first_thread_value, value, sizeof(first_thread_value)); + } + }); + std::thread second_fetch([&]() { + while (!first_fetch_ready.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + const void *value = provider->get_vector(key_for_id(47)); + if (value != nullptr) { + std::memcpy(&second_thread_value, value, sizeof(second_thread_value)); + } + second_fetch_done.store(true, std::memory_order_release); + }); + + constexpr auto kConcurrentFetchTimeout = std::chrono::seconds(5); + const auto concurrent_fetch_deadline = + std::chrono::steady_clock::now() + kConcurrentFetchTimeout; + while (!second_fetch_done.load(std::memory_order_acquire) && + std::chrono::steady_clock::now() < concurrent_fetch_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + const bool concurrent_fetch_completed = + second_fetch_done.load(std::memory_order_acquire); + // Let the first worker exit before joining even when the second fetch is + // starved. On Windows this releases its IOCP association, so a regression is + // reported as a bounded test failure rather than hanging the whole suite. + release_first_fetch.store(true, std::memory_order_release); + first_fetch.join(); + second_fetch.join(); + EXPECT_TRUE(concurrent_fetch_completed); + EXPECT_EQ(31.0f, first_thread_value); + EXPECT_EQ(47.0f, second_thread_value); + // The provider owns one heavyweight fetch context regardless of how many + // transient worker threads call get_vector(). Only result bytes are local + // to each thread. + EXPECT_EQ(1U, tracking_reader->context_count()); + + // Providers also need independent result bytes on the same thread. Keep the + // first provider's pointer live while a second provider fetches a different + // vector; a function-level TLS string shared by all providers overwrites it. + const void *second_provider_vector = + second_provider->get_vector(key_for_id(63)); + ASSERT_NE(second_provider_vector, nullptr); + float second_provider_value = -1.0f; + std::memcpy(&second_provider_value, second_provider_vector, + sizeof(second_provider_value)); + EXPECT_EQ(63.0f, second_provider_value); + + float retained_provider_value = -1.0f; + std::memcpy(&retained_provider_value, provider_vector, + sizeof(retained_provider_value)); + EXPECT_EQ(17.0f, retained_provider_value); + // Each provider still owns exactly one heavyweight fetch context, rather + // than retaining one for every historical worker thread. + EXPECT_EQ(2U, tracking_reader->context_count()); + + second_provider.reset(); + provider.reset(); + + const void *iterator_vector = provider_iterator->data(); + ASSERT_NE(iterator_vector, nullptr); + auto *iterator_fetch_context = + DiskAnnProviderTestPeer::iterator_context(provider_iterator.get()); + ASSERT_NE(iterator_fetch_context, nullptr); + EXPECT_EQ(expected_fetch_buffer_size(*iterator_fetch_context), + iterator_fetch_context->sector_buffer_size()); + std::memcpy(&iterator_value, iterator_vector, sizeof(iterator_value)); + EXPECT_EQ(0.0f, iterator_value); + provider_iterator->next(); + ASSERT_TRUE(provider_iterator->is_valid()); + EXPECT_EQ(key_for_id(1), provider_iterator->key()); + iterator_vector = provider_iterator->data(); + ASSERT_NE(iterator_vector, nullptr); + std::memcpy(&iterator_value, iterator_vector, sizeof(iterator_value)); + EXPECT_EQ(1.0f, iterator_value); + + provider_iterator.reset(); + streamer.reset(); + // Cached nodes keep their coordinates and adjacency lists in separate // buffers. Fetching a cached vector must read the coordinate cache rather // than returning bytes from the neighbor cache. @@ -867,12 +1592,16 @@ TEST_F(DiskAnnSearcherTest, TestFetchVector) { } ASSERT_EQ(0, cached_searcher->unload()); - ASSERT_EQ(0, ::truncate(path.c_str(), 0)); - std::string vector_after_truncate; + auto *diskann_searcher = dynamic_cast(searcher.get()); + ASSERT_NE(diskann_searcher, nullptr); + DiskAnnCacheTestPeer::set_reader( + diskann_searcher, std::make_shared( + DiskAnnCacheTestPeer::reader(diskann_searcher))); + std::string vector_after_failure; EXPECT_EQ(IndexError_Runtime, searcher->get_vector(key_for_id(doc_cnt - 1), linearCtx, - vector_after_truncate)); - EXPECT_TRUE(vector_after_truncate.empty()); + vector_after_failure)); + EXPECT_TRUE(vector_after_failure.empty()); } TEST_F(DiskAnnSearcherTest, TestFp16Entrypoint) { diff --git a/tests/core/framework/index_meta_test.cc b/tests/core/framework/index_meta_test.cc new file mode 100644 index 000000000..3dfd9eef2 --- /dev/null +++ b/tests/core/framework/index_meta_test.cc @@ -0,0 +1,103 @@ +// 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 + +namespace zvec::core { +namespace { + +void WriteHeaderWord(std::string *buffer, size_t word, uint32_t value) { + ASSERT_NE(nullptr, buffer); + ASSERT_GE(buffer->size(), (word + 1) * sizeof(value)); + std::memcpy(buffer->data() + word * sizeof(value), &value, sizeof(value)); +} + +TEST(IndexMeta, RejectsMalformedHeaderAndAttachmentBounds) { + IndexMeta source(IndexMeta::DataType::DT_FP32, 8); + source.set_trainer("test_trainer", 1, ailego::Params()); + std::string serialized; + source.serialize(&serialized); + ASSERT_FALSE(serialized.empty()); + + IndexMeta parsed; + EXPECT_FALSE(parsed.deserialize(nullptr, serialized.size())); + EXPECT_FALSE(parsed.deserialize(serialized.data(), sizeof(uint32_t))); + + std::string oversized_header = serialized; + WriteHeaderWord(&oversized_header, 0, + static_cast(serialized.size() + 1)); + EXPECT_FALSE( + parsed.deserialize(oversized_header.data(), oversized_header.size())); + + std::string wrapped_attachment = serialized; + WriteHeaderWord(&wrapped_attachment, 7, + std::numeric_limits::max() - 3); + WriteHeaderWord(&wrapped_attachment, 8, 8); + EXPECT_FALSE( + parsed.deserialize(wrapped_attachment.data(), wrapped_attachment.size())); + + std::string attachment_in_header = serialized; + WriteHeaderWord(&attachment_in_header, 7, sizeof(uint32_t)); + EXPECT_FALSE(parsed.deserialize(attachment_in_header.data(), + attachment_in_header.size())); +} + +TEST(IndexMeta, RejectsInvalidTypesUnitsAndElementSizeOverflow) { + IndexMeta source(IndexMeta::DataType::DT_FP32, 8); + std::string serialized; + source.serialize(&serialized); + + IndexMeta parsed; + std::string invalid_type = serialized; + WriteHeaderWord(&invalid_type, 3, + static_cast(IndexMeta::DataType::DT_BINARY64) + 1); + EXPECT_FALSE(parsed.deserialize(invalid_type.data(), invalid_type.size())); + + std::string invalid_unit = serialized; + WriteHeaderWord(&invalid_unit, 5, sizeof(double)); + EXPECT_FALSE(parsed.deserialize(invalid_unit.data(), invalid_unit.size())); + + std::string overflowing_size = serialized; + WriteHeaderWord(&overflowing_size, 3, + static_cast(IndexMeta::DataType::DT_FP64)); + WriteHeaderWord(&overflowing_size, 4, std::numeric_limits::max()); + WriteHeaderWord(&overflowing_size, 5, sizeof(double)); + EXPECT_FALSE( + parsed.deserialize(overflowing_size.data(), overflowing_size.size())); +} + +TEST(IndexMeta, DeserializesFromUnalignedStorage) { + IndexMeta source(IndexMeta::DataType::DT_FP16, 7); + source.set_metric("SquaredEuclidean", 3, ailego::Params()); + std::string serialized; + source.serialize(&serialized); + + std::string unaligned(1, '\0'); + unaligned.append(serialized); + + IndexMeta parsed; + ASSERT_TRUE(parsed.deserialize(unaligned.data() + 1, serialized.size())); + EXPECT_EQ(source.data_type(), parsed.data_type()); + EXPECT_EQ(source.dimension(), parsed.dimension()); + EXPECT_EQ(source.element_size(), parsed.element_size()); + EXPECT_EQ(source.metric_name(), parsed.metric_name()); +} + +} // namespace +} // namespace zvec::core diff --git a/tests/core/interface/index_group_by_test.cc b/tests/core/interface/index_group_by_test.cc index 0c7cfe877..a34410f05 100644 --- a/tests/core/interface/index_group_by_test.cc +++ b/tests/core/interface/index_group_by_test.cc @@ -277,7 +277,12 @@ class GroupByInterfaceTest : public ::testing::Test { ASSERT_EQ( 0, index->open(index_name, {StorageOptions::StorageType::kMMAP, true})) << test_case.name; - ASSERT_EQ(0, index->merge({source}, IndexFilter())) << test_case.name; + // Unsupported combinations are rejected by Index::search before training + // or backend-specific search setup. Do not turn this API validation into + // an integration build of every rejected backend (notably DiskAnn). + if (!expect_error) { + ASSERT_EQ(0, index->merge({source}, IndexFilter())) << test_case.name; + } auto query_param = test_case.query_param->clone(); AttachGroupBy(query_param); @@ -590,14 +595,22 @@ TEST_F(GroupByInterfaceTest, UnsupportedIndexTypes) { /*is_sparse=*/false, /*dimension=*/64}, #endif #if DISKANN_SUPPORTED - {"unsupported_diskann_graph", DenseDiskAnnParam(), DiskAnnQuery()}, + {"unsupported_diskann_graph", DenseDiskAnnParam(), DiskAnnQuery(), + /*is_sparse=*/false, /*dimension=*/kDimension, + /*with_refiner=*/false}, {"unsupported_diskann_linear", DenseDiskAnnParam(), - DiskAnnQuery(/*fetch_vector=*/false, /*is_linear=*/true)}, + DiskAnnQuery(/*fetch_vector=*/false, /*is_linear=*/true), + /*is_sparse=*/false, /*dimension=*/kDimension, + /*with_refiner=*/false}, {"unsupported_diskann_bf_pks", DenseDiskAnnParam(), DiskAnnQuery(/*fetch_vector=*/false, /*is_linear=*/false, - /*with_bf_pks=*/true)}, + /*with_bf_pks=*/true), + /*is_sparse=*/false, /*dimension=*/kDimension, + /*with_refiner=*/false}, {"unsupported_diskann_fetch_vector", DenseDiskAnnParam(), - DiskAnnQuery(/*fetch_vector=*/true)}, + DiskAnnQuery(/*fetch_vector=*/true), + /*is_sparse=*/false, /*dimension=*/kDimension, + /*with_refiner=*/false}, #endif }; diff --git a/tests/core/interface/index_interface_test.cc b/tests/core/interface/index_interface_test.cc index b70f1a21d..f39cd6d3d 100644 --- a/tests/core/interface/index_interface_test.cc +++ b/tests/core/interface/index_interface_test.cc @@ -27,10 +27,12 @@ #include "zvec/core/framework/index_provider.h" #endif #include +#include #include #include #include "algorithm/hnsw/hnsw_params.h" #include "algorithm/vamana/vamana_streamer.h" +#include "core/interface/indexes/holder_builder.h" #include "zvec/core/framework/index_error.h" #include "zvec/core/interface/index.h" #include "zvec/core/interface/index_factory.h" @@ -44,12 +46,513 @@ using namespace zvec::core_interface; +namespace { + +class HolderBuilderTestConverter final : public zvec::core::IndexConverter { + public: + explicit HolderBuilderTestConverter(int train_result) + : train_result_(train_result), + meta_(zvec::core::IndexMeta::DataType::DT_FP32, 2) {} + + int init(const zvec::core::IndexMeta &, + const zvec::ailego::Params &) override { + return 0; + } + + int cleanup() override { + return 0; + } + + int train(zvec::core::IndexHolder::Pointer) override { + return train_result_; + } + + int transform(zvec::core::IndexHolder::Pointer) override { + return 0; + } + + int dump(const zvec::core::IndexDumper::Pointer &) override { + return 0; + } + + const Stats &stats() const override { + return stats_; + } + + const zvec::core::IndexMeta &meta() const override { + return meta_; + } + + private: + int train_result_{0}; + Stats stats_{}; + zvec::core::IndexMeta meta_{}; +}; + +} // namespace + TEST(IndexInterface, IndexTypeKeepsExistingValues) { EXPECT_EQ(5, static_cast(IndexType::kDiskAnn)); EXPECT_EQ(6, static_cast(IndexType::kVamana)); EXPECT_EQ(7, static_cast(IndexType::kIVFRabitq)); } +TEST(IndexInterface, RejectsInvalidBaseVectorMetadata) { + for (const int dimension : {0, -1, MAX_DIMENSION + 1}) { + auto param = FlatIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(dimension) + .build(); + EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*param)) << dimension; + } + + auto invalid_data_type = FlatIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(static_cast( + static_cast(DataType::DT_BINARY64) + 1)) + .with_dimension(8) + .build(); + EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*invalid_data_type)); + + // Sparse vectors do not use a fixed dense dimension, so their established + // zero-dimension metadata remains valid. + auto sparse = FlatIndexParamBuilder() + .with_metric_type(MetricType::kInnerProduct) + .with_data_type(DataType::DT_FP32) + .with_is_sparse(true) + .build(); + EXPECT_NE(nullptr, IndexFactory::CreateAndInitIndex(*sparse)); +} + +TEST(IndexInterface, BuildMultiPassHolderPropagatesConverterFailures) { + std::vector values{1.0F, 2.0F}; + std::vector> doc_cache; + doc_cache.emplace_back( + 0, std::string(reinterpret_cast(values.data()), + values.size() * sizeof(float))); + + zvec::core::IndexHolder::Pointer holder; + auto failing_converter = std::make_shared( + zvec::core::IndexError_ReadData); + EXPECT_EQ(zvec::core::IndexError_ReadData, + BuildMultiPassHolder(DataType::DT_FP32, 2, doc_cache, + failing_converter, &holder)); + + auto empty_result_converter = std::make_shared(0); + EXPECT_EQ(zvec::core::IndexError_Runtime, + BuildMultiPassHolder(DataType::DT_FP32, 2, doc_cache, + empty_result_converter, &holder)); +} + +TEST(IndexInterface, ConverterCleanupReleasesTransientResult) { + std::vector values{1.0F, 2.0F}; + std::vector> doc_cache; + doc_cache.emplace_back( + 0, std::string(reinterpret_cast(values.data()), + values.size() * sizeof(float))); + + for (const char *converter_name : + {"HalfFloatConverter", "CosineNormalizeConverter", + "Int8StreamingConverter", "UniformUint7Converter", + "UniformUint8Converter"}) { + SCOPED_TRACE(converter_name); + zvec::core::IndexMeta meta(zvec::core::IndexMeta::DataType::DT_FP32, 2); + meta.set_metric("SquaredEuclidean", 0, zvec::ailego::Params()); + auto converter = zvec::core::IndexFactory::CreateConverter(converter_name); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(meta, zvec::ailego::Params())); + + zvec::core::IndexHolder::Pointer holder; + ASSERT_EQ(0, BuildMultiPassHolder(DataType::DT_FP32, 2, doc_cache, + converter, &holder)); + ASSERT_NE(nullptr, converter->result()); + ASSERT_EQ(0, converter->cleanup()); + EXPECT_EQ(nullptr, converter->result()); + + // Releasing the transient result must not make a converter single-use; + // DiskAnn retrains it during a later merge. + holder.reset(); + ASSERT_EQ(0, BuildMultiPassHolder(DataType::DT_FP32, 2, doc_cache, + converter, &holder)); + ASSERT_NE(nullptr, converter->result()); + EXPECT_EQ(0, converter->cleanup()); + } +} + +TEST(IndexInterface, DiskAnnParamJsonRoundTrip) { + auto param = DiskAnnIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(768) + .with_max_degree(48) + .with_list_size(80) + .with_pq_chunk_num(16) + .build(); + + auto restored = + IndexFactory::DeserializeIndexParamFromJson(param->serialize_to_json()); + auto diskann = std::dynamic_pointer_cast(restored); + ASSERT_NE(nullptr, diskann); + EXPECT_EQ(48, diskann->max_degree); + EXPECT_EQ(80, diskann->list_size); + EXPECT_EQ(16, diskann->pq_chunk_num); + + auto invalid = DiskAnnIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(64) + .with_max_degree(-1) + .with_list_size(32) + .with_pq_chunk_num(8) + .build(); + EXPECT_EQ(nullptr, IndexFactory::DeserializeIndexParamFromJson( + invalid->serialize_to_json())); + invalid->max_degree = 32; + invalid->list_size = -1; + EXPECT_EQ(nullptr, IndexFactory::DeserializeIndexParamFromJson( + invalid->serialize_to_json())); + invalid->list_size = 32; + invalid->pq_chunk_num = -1; + EXPECT_EQ(nullptr, IndexFactory::DeserializeIndexParamFromJson( + invalid->serialize_to_json())); +#if DISKANN_SUPPORTED + EXPECT_EQ(nullptr, IndexFactory::CreateAndInitIndex(*invalid)); +#endif +} + +TEST(IndexInterface, DiskAnnQueryParamJsonRoundTrip) { + DiskAnnQueryParam param; + param.topk = 12; + param.fetch_vector = true; + param.radius = 0.25F; + param.is_linear = true; + param.list_size = 321; + + const std::string json = IndexFactory::QueryParamSerializeToJson(param); + auto typed = + IndexFactory::QueryParamDeserializeFromJson(json); + ASSERT_NE(nullptr, typed); + EXPECT_EQ(param.topk, typed->topk); + EXPECT_EQ(param.fetch_vector, typed->fetch_vector); + EXPECT_FLOAT_EQ(param.radius, typed->radius); + EXPECT_EQ(param.is_linear, typed->is_linear); + EXPECT_EQ(param.list_size, typed->list_size); + + auto polymorphic = + IndexFactory::QueryParamDeserializeFromJson(json); + auto diskann = std::dynamic_pointer_cast(polymorphic); + ASSERT_NE(nullptr, diskann); + EXPECT_EQ(param.list_size, diskann->list_size); + + EXPECT_EQ( + nullptr, + IndexFactory::QueryParamDeserializeFromJson( + R"({"index_type":"kDiskAnn","topk":1,"fetch_vector":false,"radius":0,"is_linear":false,"list_size":0})")); +} + +#if DISKANN_SUPPORTED +TEST(IndexInterface, DiskAnnFetchHandlesSparseDocumentIdsAcrossLifecycle) { + constexpr uint32_t kSparseDocId = 1'000'000'000U; + const std::string path{"diskann_fetch_untrained.index"}; + zvec::test_util::RemoveTestFiles(path); + + auto param = DiskAnnIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(4) + .with_max_degree(16) + .with_list_size(32) + .with_pq_chunk_num(2) + .build(); + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + EXPECT_EQ(0U, index->get_doc_count()); + + std::array first_vector{1.0F, 2.0F, 3.0F, 4.0F}; + EXPECT_EQ(zvec::core::IndexError_NoReady, + index->add(VectorData{DenseVector{first_vector.data()}}, 2)); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kMMAP, true})); + EXPECT_EQ(0U, index->get_doc_count()); + + std::array second_vector{5.0F, 6.0F, 7.0F, 8.0F}; + ASSERT_EQ(0, index->add(VectorData{DenseVector{first_vector.data()}}, 2)); + ASSERT_EQ(0, index->add(VectorData{DenseVector{second_vector.data()}}, + kSparseDocId)); + EXPECT_EQ(2U, index->get_doc_count()); + EXPECT_EQ(zvec::core::IndexError_InvalidArgument, + index->add(VectorData{DenseVector{nullptr}}, 8)); + EXPECT_EQ(zvec::core::IndexError_OutOfRange, + index->add(VectorData{DenseVector{first_vector.data()}}, + (std::numeric_limits::max)())); + + VectorDataBuffer fetched; + EXPECT_EQ(zvec::core::IndexError_NoExist, index->fetch(0, &fetched)); + EXPECT_EQ(zvec::core::IndexError_NoExist, index->fetch(99, &fetched)); + EXPECT_EQ(zvec::core::IndexError_InvalidArgument, index->fetch(2, nullptr)); + EXPECT_EQ(0, index->fetch(2, &fetched)); + const auto &dense = std::get(fetched.vector_buffer); + ASSERT_EQ(sizeof(first_vector), dense.data.size()); + EXPECT_EQ(0, std::memcmp(first_vector.data(), dense.data.data(), + sizeof(first_vector))); + + ASSERT_EQ(0, index->train()); + EXPECT_EQ(2U, index->get_doc_count()); + EXPECT_EQ(zvec::core::IndexError_NoExist, index->fetch(0, &fetched)); + EXPECT_EQ(zvec::core::IndexError_NoExist, index->fetch(1, &fetched)); + EXPECT_EQ(0, index->fetch(2, &fetched)); + EXPECT_EQ(0, + std::memcmp( + first_vector.data(), + std::get(fetched.vector_buffer).data.data(), + sizeof(first_vector))); + EXPECT_EQ(0, index->fetch(kSparseDocId, &fetched)); + EXPECT_EQ(0, + std::memcmp( + second_vector.data(), + std::get(fetched.vector_buffer).data.data(), + sizeof(second_vector))); + + EXPECT_EQ(0, index->close()); + EXPECT_EQ(0U, index->get_doc_count()); + + auto reopened = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, reopened); + ASSERT_EQ(0, + reopened->open(path, {StorageOptions::StorageType::kMMAP, + /*create_new=*/false, /*read_only=*/true})); + EXPECT_EQ(2U, reopened->get_doc_count()); + EXPECT_EQ(zvec::core::IndexError_NoExist, reopened->fetch(1, &fetched)); + EXPECT_EQ(0, reopened->fetch(kSparseDocId, &fetched)); + EXPECT_EQ(0, + std::memcmp( + second_vector.data(), + std::get(fetched.vector_buffer).data.data(), + sizeof(second_vector))); + EXPECT_EQ(0, reopened->close()); + EXPECT_EQ(0U, reopened->get_doc_count()); + + auto read_only_new = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, read_only_new); + ASSERT_EQ( + 0, read_only_new->open(path, {StorageOptions::StorageType::kMMAP, + /*create_new=*/true, /*read_only=*/true})); + EXPECT_EQ( + zvec::core::IndexError_Runtime, + read_only_new->add(VectorData{DenseVector{first_vector.data()}}, 2)); + EXPECT_EQ(0, read_only_new->close()); + + zvec::test_util::RemoveTestFiles(path); +} + +TEST(IndexInterface, DiskAnnTrainCanRetryAfterSnapshotCommitFailure) { + constexpr uint32_t kDimension = 8; + constexpr uint32_t kDocCount = 16; + const std::string path{"diskann_train_retry.index"}; + zvec::test_util::RemoveTestFiles(path); + zvec::ailego::FileHelper::RemoveDirectory(path.c_str()); + + auto param = DiskAnnIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_max_degree(8) + .with_list_size(16) + .with_pq_chunk_num(2) + .build(); + auto index = IndexFactory::CreateAndInitIndex(*param); + ASSERT_NE(nullptr, index); + ASSERT_EQ(0, index->open(path, {StorageOptions::StorageType::kMMAP, true})); + + std::vector> vectors(kDocCount); + for (uint32_t id = 0; id < kDocCount; ++id) { + for (uint32_t dim = 0; dim < kDimension; ++dim) { + vectors[id][dim] = static_cast(id * kDimension + dim); + } + ASSERT_EQ(0, index->add(VectorData{DenseVector{vectors[id].data()}}, id)); + } + + // A directory cannot be atomically replaced by the completed index file. + // Once that transient obstruction is removed, the same Index instance must + // be able to rebuild and commit without remaining in BUILT state. + ASSERT_TRUE(zvec::ailego::FileHelper::MakePath(path.c_str())); + EXPECT_NE(0, index->train()); + ASSERT_TRUE(zvec::ailego::FileHelper::RemoveDirectory(path.c_str())); + ASSERT_EQ(0, index->train()); + + VectorDataBuffer fetched; + ASSERT_EQ(0, index->fetch(7, &fetched)); + const auto &dense = std::get(fetched.vector_buffer); + ASSERT_EQ(sizeof(vectors[7]), dense.data.size()); + EXPECT_EQ( + 0, std::memcmp(vectors[7].data(), dense.data.data(), dense.data.size())); + + EXPECT_EQ(0, index->close()); + zvec::test_util::RemoveTestFiles(path); +} + +TEST(IndexInterface, DiskAnnSupportsRepeatedMerge) { + constexpr uint32_t kDimension = 8; + constexpr uint32_t kDocCount = 64; + const std::string first_path{"diskann_repeat_merge_source_1.index"}; + const std::string second_path{"diskann_repeat_merge_source_2.index"}; + const std::string target_path{"diskann_repeat_merge_target.index"}; + for (const auto &path : {first_path, second_path, target_path}) { + zvec::test_util::RemoveTestFiles(path); + } + + auto flat_param = FlatIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .build(); + auto first = IndexFactory::CreateAndInitIndex(*flat_param); + auto second = IndexFactory::CreateAndInitIndex(*flat_param); + ASSERT_NE(nullptr, first); + ASSERT_NE(nullptr, second); + ASSERT_EQ( + 0, first->open(first_path, {StorageOptions::StorageType::kMMAP, true})); + ASSERT_EQ( + 0, second->open(second_path, {StorageOptions::StorageType::kMMAP, true})); + + std::vector> first_vectors(kDocCount); + std::vector> second_vectors(kDocCount); + for (uint32_t id = 0; id < kDocCount; ++id) { + for (uint32_t dim = 0; dim < kDimension; ++dim) { + first_vectors[id][dim] = static_cast(id + dim); + second_vectors[id][dim] = static_cast(1000 + id + dim); + } + ASSERT_EQ( + 0, first->add(VectorData{DenseVector{first_vectors[id].data()}}, id)); + ASSERT_EQ( + 0, second->add(VectorData{DenseVector{second_vectors[id].data()}}, id)); + } + + auto diskann_param = + DiskAnnIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_max_degree(16) + .with_list_size(32) + .with_pq_chunk_num(2) + .with_quantizer_param(QuantizerParam(QuantizerType::kFP16)) + .build(); + auto target = IndexFactory::CreateAndInitIndex(*diskann_param); + ASSERT_NE(nullptr, target); + ASSERT_EQ( + 0, target->open(target_path, {StorageOptions::StorageType::kMMAP, true})); + ASSERT_EQ(0, target->merge({first}, IndexFilter())); + auto first_snapshot = target->index_searcher(); + ASSERT_NE(nullptr, first_snapshot); + ASSERT_EQ(0, target->merge({second}, IndexFilter())); + + // A caller that acquired the previous snapshot before the atomic swap keeps + // it alive. Commit must not unload its file reader while an in-flight search + // or provider can still use it. + auto first_snapshot_provider = first_snapshot->create_provider(); + ASSERT_NE(nullptr, first_snapshot_provider); + EXPECT_EQ(kDocCount, first_snapshot_provider->count()); + + VectorDataBuffer fetched; + ASSERT_EQ(0, target->fetch(7, &fetched)); + const auto &dense = std::get(fetched.vector_buffer); + ASSERT_EQ(sizeof(second_vectors[7]), dense.data.size()); + EXPECT_EQ(0, std::memcmp(second_vectors[7].data(), dense.data.data(), + sizeof(second_vectors[7]))); + + EXPECT_EQ(0, target->close()); + first_snapshot_provider.reset(); + first_snapshot.reset(); + + auto read_only_target = IndexFactory::CreateAndInitIndex(*diskann_param); + ASSERT_NE(nullptr, read_only_target); + ASSERT_EQ(0, read_only_target->open( + target_path, {StorageOptions::StorageType::kMMAP, + /*create_new=*/false, /*read_only=*/true})); + EXPECT_EQ(zvec::core::IndexError_Runtime, + read_only_target->merge({first}, IndexFilter())); + VectorDataBuffer read_only_fetched; + ASSERT_EQ(0, read_only_target->fetch(7, &read_only_fetched)); + const auto &read_only_dense = + std::get(read_only_fetched.vector_buffer); + ASSERT_EQ(sizeof(second_vectors[7]), read_only_dense.data.size()); + EXPECT_EQ(0, + std::memcmp(second_vectors[7].data(), read_only_dense.data.data(), + sizeof(second_vectors[7]))); + EXPECT_EQ(0, read_only_target->close()); + + EXPECT_EQ(0, first->close()); + EXPECT_EQ(0, second->close()); + for (const auto &path : {first_path, second_path, target_path}) { + zvec::test_util::RemoveTestFiles(path); + } +} + +TEST(IndexInterface, DiskAnnContextUsesCurrentIndexListSize) { + static constexpr uint32_t kDimension = 8; + constexpr uint32_t kLargeDocCount = 32; + const std::string small_path{"diskann_context_small.index"}; + const std::string large_path{"diskann_context_large.index"}; + for (const auto &path : {small_path, large_path}) { + zvec::test_util::RemoveTestFiles(path); + } + + auto diskann_param = DiskAnnIndexParamBuilder() + .with_metric_type(MetricType::kL2sq) + .with_data_type(DataType::DT_FP32) + .with_dimension(kDimension) + .with_max_degree(16) + .with_list_size(32) + .with_pq_chunk_num(2) + .build(); + + auto build_index = [&](const std::string &path, uint32_t count) { + auto index = IndexFactory::CreateAndInitIndex(*diskann_param); + EXPECT_NE(nullptr, index); + if (index == nullptr) { + return index; + } + EXPECT_EQ(0, index->open(path, {StorageOptions::StorageType::kMMAP, true})); + for (uint32_t id = 0; id < count; ++id) { + std::array vector{}; + vector.fill(static_cast(id)); + EXPECT_EQ(0, index->add(VectorData{DenseVector{vector.data()}}, id)); + } + EXPECT_EQ(0, index->train()); + return index; + }; + + auto small = build_index(small_path, 1); + auto large = build_index(large_path, kLargeDocCount); + ASSERT_NE(nullptr, small); + ASSERT_NE(nullptr, large); + + auto query_param = std::make_shared(); + query_param->topk = 1; + query_param->list_size = 64; + std::array small_query{}; + SearchResult small_result; + ASSERT_EQ(0, small->search(VectorData{DenseVector{small_query.data()}}, + query_param, &small_result)); + ASSERT_EQ(1U, small_result.doc_list_.size()); + + query_param->topk = 8; + std::array large_query{}; + large_query.fill(7.0F); + SearchResult large_result; + ASSERT_EQ(0, large->search(VectorData{DenseVector{large_query.data()}}, + query_param, &large_result)); + EXPECT_EQ(query_param->topk, large_result.doc_list_.size()); + + EXPECT_EQ(0, small->close()); + EXPECT_EQ(0, large->close()); + for (const auto &path : {small_path, large_path}) { + zvec::test_util::RemoveTestFiles(path); + } +} +#endif + #if RABITQ_SUPPORTED TEST(IndexInterface, IvfRabitqValidatesBuildParams) { auto make_param = [](int nlist, int sample_count) { diff --git a/tests/core/interface/mixed_streamer_reducer_test.cc b/tests/core/interface/mixed_streamer_reducer_test.cc new file mode 100644 index 000000000..9c31ad95a --- /dev/null +++ b/tests/core/interface/mixed_streamer_reducer_test.cc @@ -0,0 +1,682 @@ +// 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 "mixed_reducer/mixed_streamer_reducer.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "mixed_reducer/mixed_reducer_params.h" + +namespace zvec { +namespace core { +namespace { + +class TestProvider : public IndexProvider { + public: + enum class IteratorMode { kValid, kMissing, kReadFailure }; + + TestProvider(const IndexMeta &meta, size_t count, IteratorMode mode) + : TestProvider(meta, count, mode, {1.0F, 2.0F}) {} + + TestProvider(const IndexMeta &meta, size_t count, IteratorMode mode, + std::array values) + : meta_(meta), mode_(mode), values_(values), keys_(count) { + std::iota(keys_.begin(), keys_.end(), uint64_t{0}); + } + + TestProvider(const IndexMeta &meta, std::vector keys, + IteratorMode mode, std::array values) + : meta_(meta), mode_(mode), values_(values), keys_(std::move(keys)) {} + + class Iterator : public IndexHolder::Iterator { + public: + Iterator(std::vector keys, bool read_failure, + std::array values) + : keys_(std::move(keys)), + read_failure_(read_failure), + vector_{values[0], values[1]} {} + + const void *data() const override { + return read_failure_ ? nullptr : vector_; + } + + bool is_valid() const override { + return index_ < keys_.size(); + } + + uint64_t key() const override { + return keys_[index_]; + } + + void next() override { + ++index_; + } + + private: + std::vector keys_; + bool read_failure_{false}; + size_t index_{0}; + float vector_[2]{1.0F, 2.0F}; + }; + + size_t count() const override { + return keys_.size(); + } + + size_t dimension() const override { + return meta_.dimension(); + } + + IndexMeta::DataType data_type() const override { + return meta_.data_type(); + } + + size_t element_size() const override { + return meta_.element_size(); + } + + IndexHolder::Iterator::Pointer create_iterator() override { + if (mode_ == IteratorMode::kMissing) { + return nullptr; + } + return IndexHolder::Iterator::Pointer( + new Iterator(keys_, mode_ == IteratorMode::kReadFailure, values_)); + } + + const void *get_vector(uint64_t) const override { + return nullptr; + } + + const std::string &owner_class() const override { + return owner_class_; + } + + private: + IndexMeta meta_; + IteratorMode mode_{IteratorMode::kValid}; + std::array values_{}; + std::vector keys_; + std::string owner_class_{"MixedStreamerReducerTest"}; +}; + +class TestStreamer : public IndexStreamer { + public: + TestStreamer(const IndexMeta &meta, IndexProvider::Pointer provider, + IndexSparseProvider::Pointer sparse_provider = + IndexSparseProvider::Pointer()) + : meta_(meta), + provider_(std::move(provider)), + sparse_provider_(std::move(sparse_provider)) {} + + int open(IndexStorage::Pointer) override { + return 0; + } + + int flush(uint64_t) override { + return 0; + } + + int close() override { + return 0; + } + + const IndexMeta &meta() const override { + return meta_; + } + + const Stats &stats() const override { + return stats_; + } + + int cleanup() override { + ++cleanup_count_; + return 0; + } + + int dump(const IndexDumper::Pointer &) override { + return dump_result_; + } + + int add_with_id_impl(uint32_t id, const void *query, + const IndexQueryMeta &qmeta, + Context::Pointer &) override { + if (query == nullptr || qmeta.data_type() != IndexMeta::DataType::DT_FP32 || + qmeta.dimension() != 2) { + return IndexError_Mismatch; + } + if (added_vectors_.size() <= id) { + added_vectors_.resize(id + 1); + } + std::memcpy(added_vectors_[id].data(), query, sizeof(float) * 2); + return 0; + } + + Provider::Pointer create_provider() const override { + return provider_; + } + + SparseProvider::Pointer create_sparse_provider() const override { + return sparse_provider_; + } + + void set_dump_result(int result) { + dump_result_ = result; + } + + const std::vector> &added_vectors() const { + return added_vectors_; + } + + size_t cleanup_count() const { + return cleanup_count_; + } + + private: + IndexMeta meta_; + IndexProvider::Pointer provider_; + IndexSparseProvider::Pointer sparse_provider_; + Stats stats_; + int dump_result_{0}; + std::vector> added_vectors_; + size_t cleanup_count_{0}; +}; + +class OffsetReformer final : public IndexReformer { + public: + OffsetReformer(float revert_offset, float convert_offset) + : revert_offset_(revert_offset), convert_offset_(convert_offset) {} + + int init(const ailego::Params &) override { + return 0; + } + + int cleanup() override { + return 0; + } + + int load(IndexStorage::Pointer) override { + return 0; + } + + int unload() override { + return 0; + } + + int revert(const void *input, const IndexQueryMeta &, + std::string *output) const override { + return transform_with_offset(input, revert_offset_, output, nullptr); + } + + int convert(const void *input, const IndexQueryMeta &, std::string *output, + IndexQueryMeta *output_meta) const override { + return transform_with_offset(input, convert_offset_, output, output_meta); + } + + private: + static int transform_with_offset(const void *input, float offset, + std::string *output, + IndexQueryMeta *output_meta) { + if (input == nullptr || output == nullptr) { + return IndexError_InvalidArgument; + } + std::array values{}; + std::memcpy(values.data(), input, sizeof(values)); + values[0] += offset; + values[1] += offset; + output->assign(reinterpret_cast(values.data()), + sizeof(values)); + if (output_meta != nullptr) { + *output_meta = IndexQueryMeta(IndexMeta::DataType::DT_FP32, 2); + } + return 0; + } + + float revert_offset_{0}; + float convert_offset_{0}; +}; + +class FailingTestBuilder final : public IndexBuilder { + public: + explicit FailingTestBuilder(int train_result) : train_result_(train_result) {} + + int cleanup() override { + return 0; + } + + const Stats &stats() const override { + return stats_; + } + + int train(IndexHolder::Pointer) override { + ++train_count_; + return train_result_; + } + + int dump(const IndexDumper::Pointer &) override { + ++dump_count_; + return 0; + } + + size_t train_count() const { + return train_count_; + } + + size_t dump_count() const { + return dump_count_; + } + + private: + int train_result_{0}; + Stats stats_; + size_t train_count_{0}; + size_t dump_count_{0}; +}; + +class TestDumper final : public IndexDumper { + public: + int init(const ailego::Params &) override { + return 0; + } + + int cleanup() override { + return 0; + } + + int create(const std::string &) override { + return 0; + } + + int close() override { + return 0; + } + + int append(const std::string &, size_t, size_t, uint32_t) override { + return 0; + } + + size_t write(const void *, size_t length) override { + return length; + } + + uint32_t magic() const override { + return 0; + } +}; + +class TestSparseProvider : public IndexSparseProvider { + public: + TestSparseProvider(size_t count, bool read_failure) + : count_(count), read_failure_(read_failure) {} + + class Iterator : public IndexSparseHolder::Iterator { + public: + Iterator(size_t count, bool read_failure) + : count_(count), read_failure_(read_failure) {} + + bool is_valid() const override { + return index_ < count_; + } + + uint64_t key() const override { + return index_; + } + + uint32_t sparse_count() const override { + return 1; + } + + const uint32_t *sparse_indices() const override { + return read_failure_ ? nullptr : &sparse_index_; + } + + const void *sparse_data() const override { + return read_failure_ ? nullptr : &sparse_value_; + } + + void next() override { + ++index_; + } + + private: + size_t count_{0}; + bool read_failure_{false}; + size_t index_{0}; + uint32_t sparse_index_{1}; + float sparse_value_{2.0F}; + }; + + size_t count() const override { + return count_; + } + + IndexMeta::DataType data_type() const override { + return IndexMeta::DataType::DT_FP32; + } + + IndexSparseHolder::Iterator::Pointer create_iterator() override { + return IndexSparseHolder::Iterator::Pointer( + new Iterator(count_, read_failure_)); + } + + size_t total_sparse_count() const override { + return count_; + } + + int get_sparse_vector(uint64_t, uint32_t *, std::string *, + std::string *) const override { + return IndexError_NotImplemented; + } + + const std::string &owner_class() const override { + return owner_class_; + } + + private: + size_t count_{0}; + bool read_failure_{false}; + std::string owner_class_{"MixedStreamerReducerSparseTest"}; +}; + +IndexMeta MakeDenseMeta() { + IndexMeta meta; + meta.set_meta_type(IndexMeta::MetaType::MT_DENSE); + meta.set_meta(IndexMeta::DataType::DT_FP32, 2); + return meta; +} + +IndexMeta MakeSparseMeta() { + return IndexMeta(IndexMeta::MetaType::MT_SPARSE, + IndexMeta::DataType::DT_FP32); +} + +int RunReduce(const IndexProvider::Pointer &target_provider, + const IndexProvider::Pointer &source_provider) { + const IndexMeta meta = MakeDenseMeta(); + auto target = std::make_shared(meta, target_provider); + auto source = std::make_shared(meta, source_provider); + + MixedStreamerReducer reducer; + ailego::Params params; + params.set(PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, 1U); + if (reducer.init(params) != 0 || + reducer.set_target_streamer_wiht_info( + nullptr, target, nullptr, nullptr, + IndexQueryMeta(IndexMeta::MetaType::MT_DENSE, + IndexMeta::DataType::DT_FP32, 2)) != 0 || + reducer.feed_streamer_with_reformer(source, nullptr) != 0) { + return IndexError_Runtime; + } + + ailego::ThreadPool thread_pool(1, false); + reducer.set_thread_pool(&thread_pool); + return reducer.reduce(IndexFilter()); +} + +int RunSparseReduce(const IndexSparseProvider::Pointer &target_provider, + const IndexSparseProvider::Pointer &source_provider) { + const IndexMeta meta = MakeSparseMeta(); + auto target = std::make_shared(meta, IndexProvider::Pointer(), + target_provider); + auto source = std::make_shared(meta, IndexProvider::Pointer(), + source_provider); + + MixedStreamerReducer reducer; + ailego::Params params; + params.set(PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, 1U); + if (reducer.init(params) != 0 || + reducer.set_target_streamer_wiht_info( + nullptr, target, nullptr, nullptr, + IndexQueryMeta(IndexMeta::MetaType::MT_SPARSE, + IndexMeta::DataType::DT_FP32)) != 0 || + reducer.feed_streamer_with_reformer(source, nullptr) != 0) { + return IndexError_Runtime; + } + + ailego::ThreadPool thread_pool(1, false); + reducer.set_thread_pool(&thread_pool); + return reducer.reduce(IndexFilter()); +} + +TEST(MixedStreamerReducer, RejectsMissingTargetProvider) { + const IndexMeta meta = MakeDenseMeta(); + auto source = std::make_shared( + meta, 1, TestProvider::IteratorMode::kValid); + EXPECT_EQ(IndexError_Runtime, RunReduce(nullptr, source)); +} + +TEST(MixedStreamerReducer, RejectsMissingSourceProvider) { + const IndexMeta meta = MakeDenseMeta(); + auto target = std::make_shared( + meta, 0, TestProvider::IteratorMode::kValid); + EXPECT_EQ(IndexError_Runtime, RunReduce(target, nullptr)); +} + +TEST(MixedStreamerReducer, RejectsMissingSourceIterator) { + const IndexMeta meta = MakeDenseMeta(); + auto target = std::make_shared( + meta, 0, TestProvider::IteratorMode::kValid); + auto source = std::make_shared( + meta, 1, TestProvider::IteratorMode::kMissing); + EXPECT_EQ(IndexError_Runtime, RunReduce(target, source)); +} + +TEST(MixedStreamerReducer, RejectsUnreadableSourceVector) { + const IndexMeta meta = MakeDenseMeta(); + auto target = std::make_shared( + meta, 0, TestProvider::IteratorMode::kValid); + auto source = std::make_shared( + meta, 1, TestProvider::IteratorMode::kReadFailure); + EXPECT_EQ(IndexError_Runtime, RunReduce(target, source)); +} + +TEST(MixedStreamerReducer, RejectsUnreadableSourceSparseVector) { + auto target = std::make_shared(0, false); + auto source = std::make_shared(1, true); + EXPECT_EQ(IndexError_Runtime, RunSparseReduce(target, source)); +} + +TEST(MixedStreamerReducer, RejectsFirstSourceVectorLayoutMismatch) { + const IndexMeta target_meta = MakeDenseMeta(); + IndexMeta source_meta; + source_meta.set_meta_type(IndexMeta::MetaType::MT_DENSE); + source_meta.set_meta(IndexMeta::DataType::DT_FP32, 1); + + auto target = std::make_shared( + target_meta, std::make_shared( + target_meta, 0, TestProvider::IteratorMode::kValid)); + auto source = std::make_shared( + source_meta, std::make_shared( + source_meta, 1, TestProvider::IteratorMode::kValid)); + + MixedStreamerReducer reducer; + ailego::Params params; + params.set(PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, 1U); + ASSERT_EQ(0, reducer.init(params)); + ASSERT_EQ(0, reducer.set_target_streamer_wiht_info( + nullptr, target, nullptr, nullptr, + IndexQueryMeta(IndexMeta::DataType::DT_FP32, 2))); + EXPECT_EQ(IndexError_InvalidArgument, + reducer.feed_streamer_with_reformer(source, nullptr)); +} + +TEST(MixedStreamerReducer, ReencodesSourcesEvenWhenReformerNamesMatch) { + IndexMeta target_meta = MakeDenseMeta(); + target_meta.set_reformer("target", 0, ailego::Params()); + IndexMeta matching_source_meta = target_meta; + IndexMeta different_source_meta = MakeDenseMeta(); + different_source_meta.set_reformer("source", 0, ailego::Params()); + + auto target = std::make_shared( + target_meta, std::make_shared( + target_meta, 0, TestProvider::IteratorMode::kValid)); + auto matching_source = std::make_shared( + matching_source_meta, + std::make_shared(matching_source_meta, 1, + TestProvider::IteratorMode::kValid, + std::array{1.0F, 2.0F})); + auto different_source = std::make_shared( + different_source_meta, + std::make_shared(different_source_meta, 1, + TestProvider::IteratorMode::kValid, + std::array{3.0F, 4.0F})); + auto target_reformer = std::make_shared(0.0F, 10.0F); + auto matching_source_reformer = std::make_shared(5.0F, 0.0F); + auto source_reformer = std::make_shared(2.0F, 0.0F); + + MixedStreamerReducer reducer; + ailego::Params params; + params.set(PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, 1U); + ASSERT_EQ(0, reducer.init(params)); + ASSERT_EQ(0, reducer.set_target_streamer_wiht_info( + nullptr, target, nullptr, target_reformer, + IndexQueryMeta(IndexMeta::DataType::DT_FP32, 2))); + ASSERT_EQ(0, reducer.feed_streamer_with_reformer(matching_source, + matching_source_reformer)); + ASSERT_EQ(0, reducer.feed_streamer_with_reformer(different_source, + source_reformer)); + + ailego::ThreadPool thread_pool(1, false); + reducer.set_thread_pool(&thread_pool); + ASSERT_EQ(0, reducer.reduce(IndexFilter())); + + const auto &added = target->added_vectors(); + ASSERT_EQ(2U, added.size()); + EXPECT_FLOAT_EQ(16.0F, added[0][0]); + EXPECT_FLOAT_EQ(17.0F, added[0][1]); + EXPECT_FLOAT_EQ(15.0F, added[1][0]); + EXPECT_FLOAT_EQ(16.0F, added[1][1]); + + target->set_dump_result(IndexError_WriteData); + EXPECT_EQ(IndexError_WriteData, reducer.dump(std::make_shared())); + EXPECT_EQ(0, reducer.cleanup()); + EXPECT_EQ(1U, target->cleanup_count()); +} + +TEST(MixedStreamerReducer, FilterOffsetsUseSourceKeySpan) { + const IndexMeta meta = MakeDenseMeta(); + auto target = std::make_shared( + meta, std::make_shared(meta, 0, + TestProvider::IteratorMode::kValid)); + auto first_source = std::make_shared( + meta, std::make_shared(meta, std::vector{0, 2}, + TestProvider::IteratorMode::kValid, + std::array{1.0F, 2.0F})); + auto second_source = std::make_shared( + meta, std::make_shared(meta, std::vector{0}, + TestProvider::IteratorMode::kValid, + std::array{3.0F, 4.0F})); + + MixedStreamerReducer reducer; + ailego::Params params; + params.set(PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, 1U); + ASSERT_EQ(0, reducer.init(params)); + ASSERT_EQ(0, reducer.set_target_streamer_wiht_info( + nullptr, target, nullptr, nullptr, + IndexQueryMeta(IndexMeta::DataType::DT_FP32, 2))); + ASSERT_EQ(0, reducer.feed_streamer_with_reformer(first_source, nullptr)); + ASSERT_EQ(0, reducer.feed_streamer_with_reformer(second_source, nullptr)); + + std::vector filtered_ids; + IndexFilter filter; + filter.set([&filtered_ids](uint64_t id) { + filtered_ids.push_back(id); + return id == 3; + }); + ailego::ThreadPool thread_pool(1, false); + reducer.set_thread_pool(&thread_pool); + ASSERT_EQ(0, reducer.reduce(filter)); + + EXPECT_EQ((std::vector{0, 2, 3}), filtered_ids); + const auto &added = target->added_vectors(); + ASSERT_EQ(2U, added.size()); + EXPECT_FLOAT_EQ(1.0F, added[0][0]); + EXPECT_FLOAT_EQ(1.0F, added[1][0]); +} + +TEST(MixedStreamerReducer, AppendsAfterGappedTargetKeySpan) { + const IndexMeta meta = MakeDenseMeta(); + auto target = std::make_shared( + meta, std::make_shared(meta, std::vector{0, 2}, + TestProvider::IteratorMode::kValid, + std::array{1.0F, 2.0F})); + auto source = std::make_shared( + meta, std::make_shared(meta, std::vector{0}, + TestProvider::IteratorMode::kValid, + std::array{3.0F, 4.0F})); + + MixedStreamerReducer reducer; + ailego::Params params; + params.set(PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, 1U); + ASSERT_EQ(0, reducer.init(params)); + ASSERT_EQ(0, reducer.set_target_streamer_wiht_info( + nullptr, target, nullptr, nullptr, + IndexQueryMeta(IndexMeta::DataType::DT_FP32, 2))); + ASSERT_EQ(0, reducer.feed_streamer_with_reformer(source, nullptr)); + + ailego::ThreadPool thread_pool(1, false); + reducer.set_thread_pool(&thread_pool); + ASSERT_EQ(0, reducer.reduce(IndexFilter())); + + const auto &added = target->added_vectors(); + ASSERT_EQ(4U, added.size()); + EXPECT_FLOAT_EQ(3.0F, added[3][0]); + EXPECT_FLOAT_EQ(4.0F, added[3][1]); +} + +TEST(MixedStreamerReducer, CleanupWithoutTargetIsSafe) { + MixedStreamerReducer reducer; + ailego::Params params; + params.set(PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, 1U); + ASSERT_EQ(0, reducer.init(params)); + EXPECT_EQ(0, reducer.cleanup()); +} + +TEST(MixedStreamerReducer, FailedTargetBuildDoesNotEnterReducedState) { + const IndexMeta meta = MakeDenseMeta(); + auto target = std::make_shared( + meta, std::make_shared(meta, 0, + TestProvider::IteratorMode::kValid)); + auto source = std::make_shared( + meta, std::make_shared(meta, 1, + TestProvider::IteratorMode::kValid)); + auto builder = std::make_shared(IndexError_ReadData); + + MixedStreamerReducer reducer; + ailego::Params params; + params.set(PARAM_MIXED_STREAMER_REDUCER_NUM_OF_ADD_THREADS, 1U); + ASSERT_EQ(0, reducer.init(params)); + ASSERT_EQ(0, reducer.set_target_streamer_wiht_info( + builder, target, nullptr, nullptr, + IndexQueryMeta(IndexMeta::DataType::DT_FP32, 2))); + ASSERT_EQ(0, reducer.feed_streamer_with_reformer(source, nullptr)); + + ailego::ThreadPool thread_pool(1, false); + reducer.set_thread_pool(&thread_pool); + EXPECT_EQ(IndexError_ReadData, reducer.reduce(IndexFilter())); + EXPECT_EQ(1U, builder->train_count()); + + EXPECT_EQ(IndexError_NoReady, reducer.dump(std::make_shared())); + EXPECT_EQ(0U, builder->dump_count()); +} + +} // namespace +} // namespace core +} // namespace zvec diff --git a/tests/core/quantizer/converter_read_failure_test.cc b/tests/core/quantizer/converter_read_failure_test.cc new file mode 100644 index 000000000..29215c9da --- /dev/null +++ b/tests/core/quantizer/converter_read_failure_test.cc @@ -0,0 +1,176 @@ +// 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 + +namespace zvec { +namespace core { +namespace { + +class TestHolder : public IndexHolder { + public: + explicit TestHolder(bool fail_reads) : fail_reads_(fail_reads) {} + + class Iterator : public IndexHolder::Iterator { + public: + Iterator(bool fail_reads, const std::array *values) + : fail_reads_(fail_reads), values_(values) {} + + const void *data() const override { + return fail_reads_ ? nullptr : values_->data(); + } + + bool is_valid() const override { + return valid_; + } + + uint64_t key() const override { + return 0; + } + + void next() override { + valid_ = false; + } + + private: + bool fail_reads_{false}; + bool valid_{true}; + const std::array *values_{nullptr}; + }; + + size_t count() const override { + return 1; + } + + size_t dimension() const override { + return values_.size(); + } + + IndexMeta::DataType data_type() const override { + return IndexMeta::DataType::DT_FP32; + } + + size_t element_size() const override { + return values_.size() * sizeof(float); + } + + bool multipass() const override { + return true; + } + + Iterator::Pointer create_iterator() override { + return Iterator::Pointer(new Iterator(fail_reads_, &values_)); + } + + private: + bool fail_reads_{false}; + std::array values_{1.0F, 2.0F, 3.0F, 4.0F}; +}; + +IndexMeta MakeMeta(const std::string &converter_name) { + IndexMeta meta(IndexMeta::DataType::DT_FP32, 4); + meta.set_metric( + converter_name == "BinaryConverter" ? "InnerProduct" : "SquaredEuclidean", + 0, ailego::Params()); + return meta; +} + +TEST(ConverterReadFailure, TransformWrappersPropagateNullData) { + for (const char *name : + {"HalfFloatConverter", "CosineNormalizeConverter", + "Int8QuantizerConverter", "Int8StreamingConverter", "MipsConverter", + "UniformUint7Converter", "UniformUint8Converter", "BinaryConverter"}) { + SCOPED_TRACE(name); + auto converter = IndexFactory::CreateConverter(name); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(MakeMeta(name), ailego::Params())); + if (std::string(name) == "Int8QuantizerConverter" || + std::string(name) == "MipsConverter" || + std::string(name) == "UniformUint7Converter" || + std::string(name) == "UniformUint8Converter") { + ASSERT_EQ(0, converter->train(std::make_shared(false))); + } + ASSERT_EQ(0, converter->transform(std::make_shared(true))); + auto result = converter->result(); + ASSERT_NE(nullptr, result); + auto iterator = result->create_iterator(); + ASSERT_NE(nullptr, iterator); + ASSERT_TRUE(iterator->is_valid()); + EXPECT_EQ(nullptr, iterator->data()); + + ASSERT_EQ(0, converter->transform(std::make_shared(false))); + iterator = converter->result()->create_iterator(); + ASSERT_NE(nullptr, iterator); + ASSERT_TRUE(iterator->is_valid()); + EXPECT_NE(nullptr, iterator->data()); + } +} + +TEST(ConverterReadFailure, TrainingReturnsReadData) { + for (const char *name : {"Int8QuantizerConverter", "MipsConverter", + "UniformUint7Converter", "UniformUint8Converter"}) { + SCOPED_TRACE(name); + auto converter = IndexFactory::CreateConverter(name); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(MakeMeta(name), ailego::Params())); + EXPECT_EQ(IndexError_ReadData, + converter->train(std::make_shared(true))); + } +} + +TEST(ConverterReadFailure, BinaryConverterInitializesQuantizer) { + auto converter = IndexFactory::CreateConverter("BinaryConverter"); + ASSERT_NE(nullptr, converter); + ASSERT_EQ(0, converter->init(MakeMeta("BinaryConverter"), ailego::Params())); + ASSERT_EQ(0, converter->transform(std::make_shared(false))); + auto iterator = converter->result()->create_iterator(); + ASSERT_NE(nullptr, iterator); + ASSERT_TRUE(iterator->is_valid()); + const void *data = iterator->data(); + ASSERT_NE(nullptr, data); + uint32_t encoded = 0; + std::memcpy(&encoded, data, sizeof(encoded)); + EXPECT_EQ(0xFU, encoded); +} + +TEST(ConverterReadFailure, BinaryReformerPreservesBatchBoundaries) { + auto reformer = IndexFactory::CreateReformer("BinaryReformer"); + ASSERT_NE(nullptr, reformer); + ASSERT_EQ(0, reformer->init(ailego::Params())); + + const std::array queries{-1.0F, -2.0F, -3.0F, -4.0F, + 1.0F, 2.0F, 3.0F, 4.0F}; + const IndexQueryMeta query_meta(IndexMeta::DataType::DT_FP32, 4); + IndexQueryMeta output_meta; + std::string output; + ASSERT_EQ(0, reformer->transform(queries.data(), query_meta, 2, &output, + &output_meta)); + ASSERT_EQ(2 * sizeof(uint32_t), output.size()); + EXPECT_EQ(IndexMeta::DataType::DT_BINARY32, output_meta.data_type()); + EXPECT_EQ(32U, output_meta.dimension()); + + std::array encoded{}; + std::memcpy(encoded.data(), output.data(), output.size()); + EXPECT_EQ(0U, encoded[0]); + EXPECT_EQ(0xFU, encoded[1]); +} + +} // namespace +} // namespace core +} // namespace zvec diff --git a/tests/core/utility/file_dumper_test.cc b/tests/core/utility/file_dumper_test.cc index 2677d4404..cab0336de 100644 --- a/tests/core/utility/file_dumper_test.cc +++ b/tests/core/utility/file_dumper_test.cc @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. #include +#include #include #include "zvec/core/framework/index_factory.h" #include "zvec/core/framework/index_helper.h" @@ -63,6 +64,50 @@ TEST(FileDumper, General) { } } +TEST(FileReadStorage, BufferedSegmentReadsOwnConstructedBytes) { + const std::string file_path = "file_read_storage_buffered_test_file"; + const std::string payload = "buffered segment payload"; + + auto dumper = IndexFactory::CreateDumper("FileDumper"); + ASSERT_NE(nullptr, dumper); + ASSERT_EQ(0, dumper->create(file_path)); + ASSERT_EQ(payload.size(), dumper->write(payload.data(), payload.size())); + ASSERT_EQ(0, dumper->append("payload", payload.size(), 0, 0)); + ASSERT_EQ(0, dumper->close()); + + auto storage = IndexFactory::CreateStorage("FileReadStorage"); + ASSERT_NE(nullptr, storage); + ASSERT_EQ(0, storage->init(ailego::Params())); + ASSERT_EQ(0, storage->open(file_path, false)); + + auto segment = storage->get("payload", 1); + ASSERT_NE(nullptr, segment); + + const void *data = nullptr; + ASSERT_EQ(payload.size(), segment->read(0, &data, payload.size())); + EXPECT_EQ(payload, + std::string(static_cast(data), payload.size())); + + IndexStorage::MemoryBlock block; + ASSERT_EQ(payload.size(), segment->read(0, block, payload.size())); + EXPECT_EQ(payload, std::string(static_cast(block.data()), + payload.size())); + + IndexStorage::SegmentData pieces[] = {{0, 8}, {9, payload.size() - 9}}; + ASSERT_TRUE(segment->read(pieces, 2)); + EXPECT_EQ( + payload.substr(0, 8), + std::string(static_cast(pieces[0].data), pieces[0].length)); + EXPECT_EQ( + payload.substr(9), + std::string(static_cast(pieces[1].data), pieces[1].length)); + + EXPECT_EQ(0U, segment->read(std::numeric_limits::max(), &data, 1)); + IndexStorage::SegmentData invalid(std::numeric_limits::max(), 1); + EXPECT_FALSE(segment->read(&invalid, 1)); + EXPECT_EQ(0, storage->close()); +} + TEST(IndexSegmentDumper, General) { std::string file_path = "index_segment_dumper_test_file"; diff --git a/tests/db/CMakeLists.txt b/tests/db/CMakeLists.txt index 975726a1e..219d979d2 100644 --- a/tests/db/CMakeLists.txt +++ b/tests/db/CMakeLists.txt @@ -23,10 +23,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 core_knn_flat core_knn_flat_sparse diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index 5616c872b..2b5e75e1f 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -3538,7 +3538,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)); @@ -6184,7 +6184,7 @@ TEST_F(CollectionTest, Feature_Optimize_IVF_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..9db80bab1 --- /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->delete_by_filter("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->add_column(added_field, "category").ok()); + ASSERT_TRUE( + collection->alter_column("category_copy", "category_renamed").ok()); + ASSERT_TRUE(collection->drop_column("category_renamed").ok()); + + auto dynamic_index = + std::make_shared(MetricType::L2, 16, 32, 2); + ASSERT_TRUE(collection->create_index(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->drop_index(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->group_by_query(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->group_by_query(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 diff --git a/tests/db/index/common/doc_test.cc b/tests/db/index/common/doc_test.cc index e47174a43..23446bddd 100644 --- a/tests/db/index/common/doc_test.cc +++ b/tests/db/index/common/doc_test.cc @@ -1465,6 +1465,41 @@ TEST(SearchQuery, ValidateAndSanitize) { EXPECT_TRUE(s.ok()) << s.message(); } + // DiskAnn list_size must be positive and the concrete parameter type must + // match the advertised index type. + { + SearchQuery query; + query.target_.field_name_ = "embedding"; + query.topk_ = 10; + std::vector query_vector(128, 1.0f); + query.target_.set_vector( + std::string(reinterpret_cast(query_vector.data()), + query_vector.size() * sizeof(float))); + FieldSchema schema = + FieldSchema("embedding", DataType::VECTOR_FP32, 128, false, + std::make_shared(MetricType::L2)); + + query.target_.query_params_ = std::make_shared(0); + auto s = query.validate(&schema, nullptr); + EXPECT_FALSE(s.ok()); + EXPECT_EQ(s.code(), StatusCode::INVALID_ARGUMENT); + + query.target_.query_params_ = std::make_shared(-1); + s = query.validate(&schema, nullptr); + EXPECT_FALSE(s.ok()); + EXPECT_EQ(s.code(), StatusCode::INVALID_ARGUMENT); + + query.target_.query_params_ = + std::make_shared(IndexType::DISKANN); + s = query.validate(&schema, nullptr); + EXPECT_FALSE(s.ok()); + EXPECT_EQ(s.code(), StatusCode::INVALID_ARGUMENT); + + query.target_.query_params_ = std::make_shared(300); + s = query.validate(&schema, nullptr); + EXPECT_TRUE(s.ok()) << s.message(); + } + // FTS clause validation { auto fts_params = std::make_shared(); diff --git a/tests/db/index/common/index_params_test.cc b/tests/db/index/common/index_params_test.cc index 9f6645aa4..654cfb5f6 100644 --- a/tests/db/index/common/index_params_test.cc +++ b/tests/db/index/common/index_params_test.cc @@ -163,6 +163,27 @@ TEST(IndexParamsTest, IVFIndexParams) { EXPECT_EQ(params.n_list(), 64); } +TEST(IndexParamsTest, DiskAnnIndexParams) { + DiskAnnIndexParams params(MetricType::L2, 48, 80, 16, QuantizeType::FP16, + QuantizerParam(true)); + + EXPECT_EQ(IndexType::DISKANN, params.type()); + EXPECT_EQ(MetricType::L2, params.metric_type()); + EXPECT_EQ(48, params.max_degree()); + EXPECT_EQ(80, params.list_size()); + EXPECT_EQ(16, params.pq_chunk_num()); + EXPECT_EQ(QuantizeType::FP16, params.quantize_type()); + EXPECT_TRUE(params.quantizer_param().enable_rotate()); + + auto cloned = params.clone(); + auto *cloned_diskann = dynamic_cast(cloned.get()); + ASSERT_NE(nullptr, cloned_diskann); + EXPECT_EQ(params, *cloned_diskann); + + cloned_diskann->set_list_size(40); + EXPECT_NE(params, *cloned_diskann); +} + #if RABITQ_SUPPORTED TEST(IndexParamsTest, IvfRabitqIndexParams) { IvfRabitqIndexParams params(MetricType::COSINE, 32, 1, 5); diff --git a/tests/db/index/common/schema_test.cc b/tests/db/index/common/schema_test.cc index 66559b272..59269d3f4 100644 --- a/tests/db/index/common/schema_test.cc +++ b/tests/db/index/common/schema_test.cc @@ -1068,6 +1068,48 @@ TEST(FieldSchemaTest, IvfRabitqIndexValidationParameters) { } #endif +#if DISKANN_SUPPORTED +TEST(FieldSchemaTest, DiskAnnRejectsUnsupportedDataAndQuantizationTypes) { + auto make_field = [](DataType data_type, QuantizeType quantize_type, + QuantizerParam quantizer_param = QuantizerParam()) { + return FieldSchema( + "vector_field", data_type, 128, false, + std::make_shared(MetricType::L2, 32, 50, 8, + quantize_type, quantizer_param)); + }; + + EXPECT_TRUE(make_field(DataType::VECTOR_FP32, QuantizeType::UNDEFINED) + .validate() + .ok()); + EXPECT_TRUE( + make_field(DataType::VECTOR_FP32, QuantizeType::FP16).validate().ok()); + EXPECT_TRUE(make_field(DataType::VECTOR_FP16, QuantizeType::UNDEFINED) + .validate() + .ok()); + + for (QuantizeType quantize_type : + {QuantizeType::INT4, QuantizeType::INT8, QuantizeType::RABITQ}) { + auto status = make_field(DataType::VECTOR_FP32, quantize_type).validate(); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.message().find("only supports FP16 quantization"), + std::string::npos); + } + + auto int8_status = + make_field(DataType::VECTOR_INT8, QuantizeType::UNDEFINED).validate(); + EXPECT_FALSE(int8_status.ok()); + EXPECT_NE(int8_status.message().find("only supports FP32/FP16"), + std::string::npos); + + auto rotate_status = make_field(DataType::VECTOR_FP32, QuantizeType::FP16, + QuantizerParam(true)) + .validate(); + EXPECT_FALSE(rotate_status.ok()); + EXPECT_NE(rotate_status.message().find("does not support quantizer rotation"), + std::string::npos); +} +#endif + TEST(FieldSchemaTest, HnswRabitqIndexValidation_UnsupportedDataTypes) { // Test unsupported data types with HNSW_RABITQ index diff --git a/tools/core/recall_original.cc b/tools/core/recall_original.cc index ce8a2fbdf..f7bb25927 100644 --- a/tools/core/recall_original.cc +++ b/tools/core/recall_original.cc @@ -12,10 +12,12 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include #include +#include +#include #include #include +#include #include #include #include @@ -67,6 +69,54 @@ enum RetrievalMode { RM_UNDEFINED = 0, RM_DENSE = 1, RM_SPARSE = 2 }; enum FilterMode { FM_UNDEFINED = 0, FM_NONE = 1, FM_TAG = 2 }; +using RecallOutputFiles = vector>; + +void close_recall_output_files(RecallOutputFiles *files) { + for (const auto &fs : *files) { + fs.first->close(); + fs.second->close(); + delete fs.first; + delete fs.second; + } + files->clear(); +} + +bool open_recall_output_files(const string &output, size_t threads, + RecallOutputFiles *files) { + if (output.empty()) { + return true; + } + + std::error_code ec; + std::filesystem::create_directories(output, ec); + if (ec) { + cerr << "Failed to create output directory [" << output + << "]: " << ec.message() << endl; + return false; + } + if (!std::filesystem::is_directory(output, ec)) { + cerr << "Invalid output directory [" << output + << "]: " << (ec ? ec.message() : "path is not a directory") << endl; + return false; + } + + cout << "logs output to : " << output << endl; + for (size_t i = 0; i < threads; ++i) { + std::unique_ptr fs_k(new fstream()); + fs_k->open(output + "/t" + to_string(i) + ".knn", ios::out); + std::unique_ptr fs_l(new fstream()); + fs_l->open(output + "/t" + to_string(i) + ".linear", ios::out); + if (!fs_k->is_open() || !fs_l->is_open()) { + cerr << "Failed to open recall output files in [" << output << "]" + << endl; + close_recall_output_files(files); + return false; + } + files->emplace_back(fs_k.release(), fs_l.release()); + } + return true; +} + template class Recall { public: @@ -103,7 +153,8 @@ class Recall { << flush; } - void run_dense(Flow *flower, const string &recall_tops, size_t gt_count) { + bool run_dense(Flow *flower, const string &recall_tops, size_t gt_count) { + worker_failed_.store(false, std::memory_order_relaxed); StringHelper::Split(recall_tops, ",", &topk_ids_); std::sort(topk_ids_.begin(), topk_ids_.end()); @@ -123,7 +174,7 @@ class Recall { if (!load_gt_dense(flower, gt_count)) { cerr << "Load ground truth file failed!" << endl; - return; + return false; } } @@ -136,25 +187,9 @@ class Recall { } // Prepare file handler - vector> output_fs; - if (!output_.empty()) { - string cmd = "mkdir -p " + output_; - int ret = system(cmd.c_str()); - if (ret != 0) { - std::cerr << "execute cmd " << cmd << " failed" << std::endl; - return; - } - struct stat sb; - if (stat(output_.c_str(), &sb) == 0 && S_ISDIR(sb.st_mode)) { - cout << "logs output to : " << output_ << endl; - for (size_t i = 0; i < threads_; ++i) { - fstream *fs_k = new fstream(); - fs_k->open(output_ + "/t" + to_string(i) + ".knn", ios::out); - fstream *fs_l = new fstream(); - fs_l->open(output_ + "/t" + to_string(i) + ".linear", ios::out); - output_fs.push_back(make_pair(fs_k, fs_l)); - } - } + RecallOutputFiles output_fs; + if (!open_recall_output_files(output_, threads_, &output_fs)) { + return false; } signal(SIGINT, stop); @@ -173,17 +208,22 @@ class Recall { } pool_->wait_finish(); - for (auto fs : output_fs) { - fs.first->close(); - fs.second->close(); - delete fs.first; - delete fs.second; + close_recall_output_files(&output_fs); + if (worker_failed_.load(std::memory_order_relaxed)) { + cerr << "Recall failed because one or more query tasks failed" << endl; + return false; + } + if (i != batch_queries_.size()) { + cerr << "Recall interrupted before all query tasks were submitted" + << endl; + return false; } cout << "Process query: " << i << endl; for (auto it : recall_res_) { cout << "Recall@" << it.first << ": " << it.second / linear_queries_.size() << endl; } + return true; } bool load_query(const std::string &query_file, const std::string &first_sep, @@ -368,13 +408,15 @@ class Recall { Flow::Context::Pointer context = flower->create_context(); if (!context) { cerr << "Failed to create search context" << endl; + error.store(true, std::memory_order_relaxed); return; } FilterResultCache filter_cache; if (filter_mode_ == FM_TAG) { - if (batch_taglists_[i].size() != 1) { + if (i >= batch_taglists_.size() || batch_taglists_[i].size() != 1) { cerr << "query tag list not equal to one!" << endl; + error.store(true, std::memory_order_relaxed); return; } @@ -383,7 +425,7 @@ class Recall { flower->tag_key_list()); if (ret != 0) { cerr << "prefilter failed, idx: " << i << std::endl; - + error.store(true, std::memory_order_relaxed); return; } @@ -682,6 +724,9 @@ class Recall { void recall_one_dense( Flow *flower, size_t topk, size_t index, std::vector> &output_fs) { + if (worker_failed_.load(std::memory_order_relaxed)) { + return; + } const auto &query = batch_queries_[index]; size_t thread_index = pool_->indexof_this(); @@ -695,6 +740,7 @@ class Recall { Flow::Context::Pointer knn_context = flower->create_context(); if (!knn_context) { cerr << "Failed to create search context" << endl; + worker_failed_.store(true, std::memory_order_relaxed); return; } knn_context->set_topk(topk); @@ -805,8 +851,10 @@ class Recall { // prefilter FilterResultCache filter_cache; if (filter_mode_ == FM_TAG) { - if (batch_taglists_[index].size() != 1) { + if (index >= batch_taglists_.size() || + batch_taglists_[index].size() != 1) { cerr << "query tag list not equal to one!" << endl; + worker_failed_.store(true, std::memory_order_relaxed); return; } @@ -815,7 +863,7 @@ class Recall { flower->tag_key_list()); if (ret != 0) { cerr << "prefilter failed, idx: " << index << std::endl; - + worker_failed_.store(true, std::memory_order_relaxed); return; } @@ -830,6 +878,7 @@ class Recall { if (ret < 0) { cerr << "Failed to knn_search batch, ret=" << ret << " " << IndexError::What(ret) << endl; + worker_failed_.store(true, std::memory_order_relaxed); return; } for (size_t i = 0; i < qnum; ++i) { @@ -846,6 +895,7 @@ class Recall { if (ret < 0) { cerr << "Failed to knn_search, ret=" << ret << " " << IndexError::What(ret) << endl; + worker_failed_.store(true, std::memory_order_relaxed); return; } auto &knn_res = knn_context->result(); @@ -886,6 +936,7 @@ class Recall { bool external_gt_file_enabled_{false}; FilterMode filter_mode_{FM_NONE}; + std::atomic_bool worker_failed_{false}; static bool STOP_NOW; }; @@ -962,8 +1013,9 @@ class SparseRecall { return 0; } - void run_sparse(SparseFlow *flower, const string &recall_tops, + bool run_sparse(SparseFlow *flower, const string &recall_tops, size_t gt_count) { + worker_failed_.store(false, std::memory_order_relaxed); StringHelper::Split(recall_tops, ",", &topk_ids_); std::sort(topk_ids_.begin(), topk_ids_.end()); @@ -983,7 +1035,7 @@ class SparseRecall { if (!load_gt_sparse(flower, gt_count)) { cerr << "Load ground truth file failed!" << endl; - return; + return false; } } @@ -996,25 +1048,9 @@ class SparseRecall { } // Prepare file handler - vector> output_fs; - if (!output_.empty()) { - string cmd = "mkdir -p " + output_; - int ret = system(cmd.c_str()); - if (ret != 0) { - std::cerr << "execute cmd " << cmd << " failed" << std::endl; - return; - } - struct stat sb; - if (stat(output_.c_str(), &sb) == 0 && S_ISDIR(sb.st_mode)) { - cout << "logs output to : " << output_ << endl; - for (size_t i = 0; i < threads_; ++i) { - fstream *fs_k = new fstream(); - fs_k->open(output_ + "/t" + to_string(i) + ".knn", ios::out); - fstream *fs_l = new fstream(); - fs_l->open(output_ + "/t" + to_string(i) + ".linear", ios::out); - output_fs.push_back(make_pair(fs_k, fs_l)); - } - } + RecallOutputFiles output_fs; + if (!open_recall_output_files(output_, threads_, &output_fs)) { + return false; } signal(SIGINT, stop); @@ -1033,17 +1069,22 @@ class SparseRecall { } pool_->wait_finish(); - for (auto fs : output_fs) { - fs.first->close(); - fs.second->close(); - delete fs.first; - delete fs.second; + close_recall_output_files(&output_fs); + if (worker_failed_.load(std::memory_order_relaxed)) { + cerr << "Recall failed because one or more query tasks failed" << endl; + return false; + } + if (i != batch_sparse_counts_.size()) { + cerr << "Recall interrupted before all query tasks were submitted" + << endl; + return false; } cout << "Process query: " << i << endl; for (auto it : recall_res_) { cout << "Recall@" << it.first << ": " << it.second / linear_queries_.size() << endl; } + return true; } bool load_query(const std::string &query_file, const std::string &first_sep, @@ -1151,6 +1192,7 @@ class SparseRecall { SparseFlow::Context::Pointer context = flower->create_context(); if (!context) { cerr << "Failed to create search context" << endl; + error.store(true, std::memory_order_relaxed); return; } @@ -1160,8 +1202,9 @@ class SparseRecall { // prefilter FilterResultCache filter_cache; if (filter_mode_ == FM_TAG) { - if (batch_taglists_[i].size() != 1) { + if (i >= batch_taglists_.size() || batch_taglists_[i].size() != 1) { cerr << "query tag list not equal to one!" << endl; + error.store(true, std::memory_order_relaxed); return; } @@ -1170,7 +1213,7 @@ class SparseRecall { flower->tag_key_list()); if (ret != 0) { cerr << "prefilter failed, idx: " << i << std::endl; - + error.store(true, std::memory_order_relaxed); return; } @@ -1394,6 +1437,9 @@ class SparseRecall { void recall_one_sparse( SparseFlow *flower, size_t topk, size_t index, std::vector> &output_fs) { + if (worker_failed_.load(std::memory_order_relaxed)) { + return; + } const auto &sparse_count = batch_sparse_counts_[index]; const auto &sparse_index = batch_sparse_indices_[index]; const auto &sparse_feature = batch_sparse_features_[index]; @@ -1409,6 +1455,7 @@ class SparseRecall { SparseFlow::Context::Pointer knn_context = flower->create_context(); if (!knn_context) { cerr << "Failed to create search context" << endl; + worker_failed_.store(true, std::memory_order_relaxed); return; } knn_context->set_topk(topk); @@ -1522,8 +1569,10 @@ class SparseRecall { FilterResultCache filter_cache; if (filter_mode_ == FM_TAG) { - if (batch_taglists_[index].size() != 1) { + if (index >= batch_taglists_.size() || + batch_taglists_[index].size() != 1) { cerr << "query tag list not equal to one!" << endl; + worker_failed_.store(true, std::memory_order_relaxed); return; } @@ -1532,7 +1581,7 @@ class SparseRecall { flower->tag_key_list()); if (ret != 0) { cerr << "prefilter failed, idx: " << index << std::endl; - + worker_failed_.store(true, std::memory_order_relaxed); return; } @@ -1542,6 +1591,8 @@ class SparseRecall { } if (call_batch_api_) { + cerr << "Sparse batch recall is not supported" << endl; + worker_failed_.store(true, std::memory_order_relaxed); // size_t qnum = sparse_count.size() / dim_; // int ret = do_knn_search(flower, knn_context, sparse_count, // sparse_index, sparse_feature, qnum); if (ret < 0) { @@ -1565,6 +1616,7 @@ class SparseRecall { if (ret < 0) { cerr << "Failed to sparse_knn_search, ret=" << ret << " " << IndexError::What(ret) << endl; + worker_failed_.store(true, std::memory_order_relaxed); return; } auto &knn_res = knn_context->result(); @@ -1610,6 +1662,7 @@ class SparseRecall { bool external_gt_file_enabled_{false}; FilterMode filter_mode_{FM_NONE}; + std::atomic_bool worker_failed_{false}; static bool STOP_NOW; }; @@ -1712,9 +1765,8 @@ int recall_dense(std::string &query_type, size_t thread_count, } } - if (load_index(flower, index_dir)) { - recall.run_dense(&flower, top_k, gt_count); - } else { + if (!load_index(flower, index_dir) || + !recall.run_dense(&flower, top_k, gt_count)) { return -1; } } else if (query_type == "int8") { @@ -1731,9 +1783,8 @@ int recall_dense(std::string &query_type, size_t thread_count, } } - if (load_index(flower, index_dir)) { - recall.run_dense(&flower, top_k, gt_count); - } else { + if (!load_index(flower, index_dir) || + !recall.run_dense(&flower, top_k, gt_count)) { return -1; } } else if (query_type == "binary") { @@ -1750,9 +1801,8 @@ int recall_dense(std::string &query_type, size_t thread_count, } } - if (load_index(flower, index_dir)) { - recall.run_dense(&flower, top_k, gt_count); - } else { + if (!load_index(flower, index_dir) || + !recall.run_dense(&flower, top_k, gt_count)) { return -1; } } else if (query_type == "binary64") { @@ -1769,13 +1819,13 @@ int recall_dense(std::string &query_type, size_t thread_count, } } - if (load_index(flower, index_dir)) { - recall.run_dense(&flower, top_k, gt_count); - } else { + if (!load_index(flower, index_dir) || + !recall.run_dense(&flower, top_k, gt_count)) { return -1; } } else { cerr << "Can not recognize type: " << query_type << endl; + return -1; } return 0; @@ -1812,13 +1862,13 @@ int recall_sparse(std::string &query_type, size_t thread_count, } } - if (load_sparse_index(flower, index_dir)) { - recall.run_sparse(&flower, top_k, gt_count); - } else { + if (!load_sparse_index(flower, index_dir) || + !recall.run_sparse(&flower, top_k, gt_count)) { return -1; } } else { cerr << "Can not recognize type: " << query_type << endl; + return -1; } return 0; @@ -2020,12 +2070,15 @@ int main(int argc, char *argv[]) { } string index_dir = config_common["IndexPath"].as(); - recall_sparse(query_type, thread_count, batch_count, top_k, gt_count, - query_file, first_sep, second_sep, ground_truth_file, - ground_truth_first_sep, ground_truth_second_sep, flower, - index_dir, log_dir, filter_mode); + ret = recall_sparse(query_type, thread_count, batch_count, top_k, gt_count, + query_file, first_sep, second_sep, ground_truth_file, + ground_truth_first_sep, ground_truth_second_sep, flower, + index_dir, log_dir, filter_mode); flower.unload(); + if (ret != 0) { + return ret; + } cout << "Recall done." << endl; } else { @@ -2077,10 +2130,10 @@ int main(int argc, char *argv[]) { string index_dir = config_common["IndexPath"].as(); if (retrieval_mode == RM_DENSE) { - recall_dense(query_type, thread_count, batch_count, top_k, gt_count, - query_file, first_sep, second_sep, ground_truth_file, - ground_truth_first_sep, ground_truth_second_sep, flower, - index_dir, log_dir, filter_mode); + ret = recall_dense(query_type, thread_count, batch_count, top_k, gt_count, + query_file, first_sep, second_sep, ground_truth_file, + ground_truth_first_sep, ground_truth_second_sep, + flower, index_dir, log_dir, filter_mode); } else { std::string mode = retrieval_mode == 1 ? "Dense" : "Sparse"; cerr << "unsupported retrieval mode: " << mode << endl; @@ -2090,6 +2143,9 @@ int main(int argc, char *argv[]) { // Cleanup flower.unload(); + if (ret != 0) { + return ret; + } cout << "Recall done." << endl; }